Update train.py

train.py

@@ -197,7 +197,12 @@ def fine_tuning_step(data_iter, colorizer, discriminator, gen_optimizer, disc_op
     for cur_disc_step in range(5):
         discriminator.zero_grad()
         
-        bw, dfm, color_for_real = next(data_iter)
+        try:
+            bw, dfm, color_for_real = next(data_iter)
+        except StopIteration:
+            data_iter = iter(ft_dataloader)
+            bw, dfm, color_for_real = next(data_iter)
+            
         bw, dfm, color_for_real = bw.to(device), dfm.to(device), color_for_real.to(device)
         
         y_real = torch.full((bw.size(0), 1), 0.9, device = device)
@@ -227,7 +232,12 @@ def fine_tuning_step(data_iter, colorizer, discriminator, gen_optimizer, disc_op
         
     colorizer.generator.zero_grad()    
 
-    bw, dfm, _ = next(data_iter)
+    try:
+        bw, dfm, _ = next(data_iter)
+    except StopIteration:
+        data_iter = iter(ft_dataloader)
+        bw, dfm, _ = next(data_iter)
+        
     bw, dfm = bw.to(device), dfm.to(device)
     
     y_real = torch.ones((bw.size(0), 1), device = device)
@@ -243,8 +253,6 @@ def fine_tuning_step(data_iter, colorizer, discriminator, gen_optimizer, disc_op
     
     generator_loss.backward()
     gen_optimizer.step()
-    
-    
         
 def fine_tuning(colorizer, discriminator, content, dataloader, iterations, colorizer_optimizer, discriminator_optimizer, ft_dataloader, device = 'cpu'):
     colorizer.generator.train()