RuntimeError: generator raised StopIteration

[lkp520]

Dear author, when I run python train.py \ --seed 100 \ --dataset cifar10 \ --dataset-path ./dataset/cifar10 \ --checkpoints-path ./checkpoints \ --batch-size 128 \ --epochs 200 \ --lr 3e-4 \ --lr-decay-steps 1e4 \ --augment True

and it train until 49438/50000 [=======================>.] - ETA: 58s - epoch: 1 - iter: 391 - step: 391 - D loss: 1.3951 - D fake: 0.6976 - D real: 0.6975 - G loss: 8.5754 - G L1: 7.8779 - G gan: 0.6975 - accuracy: 0.0927 It stop and meet the problem:

Traceback (most recent call last):
File "/home/billy/kp_pytorch/Colorizing-with-GANs/src/dataset.py", line 73, in generator
raise StopIteration
StopIteration

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
File "train.py", line 5, in
main(options)
File "/home/billy/kp_pytorch/Colorizing-with-GANs/src/main.py", line 56, in main
model.train()
File "/home/billy/kp_pytorch/Colorizing-with-GANs/src/models.py", line 46, in train
for input_rgb in generator:
RuntimeError: generator raised StopIteration

How can I solve it? Thank you very much!