fix fast neural transfer on cuda (#392)

pytorch · Jul 26, 2018 · b5f3612 · b5f3612
1 parent 98b1fce
commit b5f3612
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/fast_neural_style/neural_style/neural_style.py b/fast_neural_style/neural_style/neural_style.py
@@ -69,7 +69,8 @@ def train(args):
             count += n_batch
             optimizer.zero_grad()
 
-            y = transformer(x.to(device))
+            x = x.to(device)
+            y = transformer(x)
 
             y = utils.normalize_batch(y)
             x = utils.normalize_batch(x)