fixed indexing in checkpoint

rlmwang · Jan 10, 2021 · e7edb01 · e7edb01
1 parent 98fc132
commit e7edb01
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchtools/train.py b/torchtools/train.py
@@ -31,7 +31,7 @@ def __init__(self, phase, metric):
     self.state_dict = None
 
   def step(self, epoch, model):
-    loss = epoch['score'][self.phase][self.metric][-1]
+    loss = epoch[self.phase][self.metric][-1]
 
     if self.best is None or loss < self.best:
       self.state_dict = model.state_dict()