Use new item() method after upgrade to pytorch 0.4.0

JohnVinyard · Jun 25, 2018 · 55d18ef · 55d18ef
1 parent 478ddc3
commit 55d18ef
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/zounds/learn/wgan.py b/zounds/learn/wgan.py
@@ -230,9 +230,9 @@ def train(self, data):
                 g_loss.backward()
                 self.generator_optim.step()
 
-                gl = g_loss.data[0]
-                dl = d_loss.data[0]
-                rl = real_mean.data[0]
+                gl = g_loss.data.item()
+                dl = d_loss.data.item()
+                rl = real_mean.data.item()
 
                 self.on_batch_complete(
                     epoch=epoch,