1 files changed, 1 insertions, 1 deletions
diff --git a/train.py b/train.py
index eeb77d2..b1a2213 100644
--- a/train.py
+++ b/train.py
@@ -58,7 +58,7 @@ trainer.extend(extensions.LogReport(trigger=trigger_log, log_name='log.txt'))
 
 if extensions.PlotReport.available():
     trainer.extend(extensions.PlotReport(
-        y_keys=['main/loss', 'test/main/loss'],
+        y_keys=['main/loss', 'test/main/loss', 'train/main/loss'],
         x_key='iteration',
         file_name='loss.png',
         trigger=trigger_log,