fix logging in dcgan example (pytorch#312)

weihuangxu · Nov 2, 2018 · 34e74e4 · 34e74e4
1 parent 4b7481e
commit 34e74e4
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/examples/gan/dcgan.py b/examples/gan/dcgan.py
@@ -308,8 +308,11 @@ def step(engine, batch):
 
     # attach running average metrics
     monitoring_metrics = ['errD', 'errG', 'D_x', 'D_G_z1', 'D_G_z2']
-    for metric in monitoring_metrics:
-        RunningAverage(alpha=alpha, output_transform=lambda x: x[metric]).attach(trainer, metric)
+    RunningAverage(alpha=alpha, output_transform=lambda x: x['errD']).attach(trainer, 'errD')
+    RunningAverage(alpha=alpha, output_transform=lambda x: x['errG']).attach(trainer, 'errG')
+    RunningAverage(alpha=alpha, output_transform=lambda x: x['D_x']).attach(trainer, 'D_x')
+    RunningAverage(alpha=alpha, output_transform=lambda x: x['D_G_z1']).attach(trainer, 'D_G_z1')
+    RunningAverage(alpha=alpha, output_transform=lambda x: x['D_G_z2']).attach(trainer, 'D_G_z2')
 
     # attach progress bar
     pbar = ProgressBar()