Update train.py

train.py

@@ -174,14 +174,15 @@ def train(colorizer, discriminator, content, dataloader, epochs, colorizer_optim
             
             
             if disc_step:
-                step_loss = discriminator_step(inputs, colorizer, discriminator, discriminator_optimizer, device)
+                step_loss = discriminator_step(inputs, colorizer, discriminator, optimizer, device)
                 sum_disc_loss += step_loss
+                disc_step  = False
             else:
-                step_loss = generator_step(inputs, colorizer, discriminator, content, generator_loss, colorizer_optimizer, device)
+                step_loss = generator_step(inputs, colorizer, discriminator, content, loss_function, optimizer, device)
                 sum_gen_loss += step_loss
+                disc_step = True
                 
-            disc_step = disc_step ^ True
-    
+            disc_step = disc_step ^ True # Toggle the disc_step variable at the end of the inner for loop
     
         print(datetime.datetime.now().time())
         print('Epoch : %d Discr loss: %.4f Gen loss : %.4f \n'%(epoch, sum_disc_loss / (n // 2 + 1), sum_gen_loss / (n // 2 + 1)))