Update train.py
#1
by
Keiser41
- opened
train.py
CHANGED
@@ -197,7 +197,7 @@ def fine_tuning_step(data_iter, colorizer, discriminator, gen_optimizer, disc_op
|
|
197 |
for cur_disc_step in range(5):
|
198 |
discriminator.zero_grad()
|
199 |
|
200 |
-
bw, dfm, color_for_real =
|
201 |
bw, dfm, color_for_real = bw.to(device), dfm.to(device), color_for_real.to(device)
|
202 |
|
203 |
y_real = torch.full((bw.size(0), 1), 0.9, device = device)
|
|
|
197 |
for cur_disc_step in range(5):
|
198 |
discriminator.zero_grad()
|
199 |
|
200 |
+
bw, dfm, color_for_real = next(data_iter)
|
201 |
bw, dfm, color_for_real = bw.to(device), dfm.to(device), color_for_real.to(device)
|
202 |
|
203 |
y_real = torch.full((bw.size(0), 1), 0.9, device = device)
|