Update train.py

train.py

@@ -246,7 +246,7 @@ def fine_tuning_step(data_iter, colorizer, discriminator, gen_optimizer, disc_op
     
     
         
-def fine_tuning(colorizer, discriminator, content, dataloader, iterations, colorizer_optimizer, discriminator_optimizer, data_iter, device = 'cpu'):
+def fine_tuning(colorizer, discriminator, content, dataloader, iterations, colorizer_optimizer, discriminator_optimizer, ft_dataloader, device = 'cpu'):
     colorizer.generator.train()
     discriminator.train()
     
@@ -265,6 +265,7 @@ def fine_tuning(colorizer, discriminator, content, dataloader, iterations, color
         disc_step = disc_step ^ True
         
         if n % 10 == 5:
+            data_iter = iter(ft_dataloader)
             fine_tuning_step(data_iter, colorizer, discriminator, colorizer_optimizer, discriminator_optimizer, device)
     
 if __name__ == '__main__':