Update train.py
Browse files
train.py
CHANGED
|
@@ -174,7 +174,7 @@ def train(colorizer, discriminator, content, dataloader, epochs, colorizer_optim
|
|
| 174 |
|
| 175 |
|
| 176 |
if disc_step:
|
| 177 |
-
step_loss = discriminator_step(inputs, colorizer, discriminator,
|
| 178 |
sum_disc_loss += step_loss
|
| 179 |
disc_step = False
|
| 180 |
else:
|
|
|
|
| 174 |
|
| 175 |
|
| 176 |
if disc_step:
|
| 177 |
+
step_loss = discriminator_step(inputs, colorizer, discriminator, gen_optimizer, device)
|
| 178 |
sum_disc_loss += step_loss
|
| 179 |
disc_step = False
|
| 180 |
else:
|