Update train.py

train.py

@@ -197,7 +197,7 @@ 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 = data_iter.next()
+        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)