support python 2 & 3, torch 0.3 & 0.4

Prinsphield · May 19, 2018 · eeba065 · eeba065
1 parent 0489aac
commit eeba065
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/ELEGANT.py b/ELEGANT.py
@@ -225,8 +225,8 @@ def save_sample_images(self, save_num=5):
 
     def save_scalar_log(self):
         scalar_info = {
-            'loss_D': self.loss_D.data[0],
-            'loss_G': self.loss_G.data[0],
+            'loss_D': self.loss_D.data.cpu().numpy(),
+            'loss_G': self.loss_G.data.cpu().numpy(),
             'G_lr'  : self.G_lr_scheduler.get_lr()[0],
             'D_lr'  : self.D_lr_scheduler.get_lr()[0],
         }
@@ -294,7 +294,7 @@ def train(self):
                 if self.step % 2000 == 0:
                     self.save_sample_images()
 
-            print('step: {:06d}, loss D: {}, loss G: {}'.format(self.step, self.loss_D.data.cpu().numpy()[0], self.loss_G.data.cpu().numpy()[0]))
+            print('step: %06d, loss D: %.6f, loss G: %.6f' % (self.step, self.loss_D.data.cpu().numpy(), self.loss_G.data.cpu().numpy()))
 
             if self.step % 100 == 0:
                 self.save_scalar_log()