Merge pull request #61 from MadryLab/fix-amp-loading-bug

Small fix regarding old non-amp saved models
MadryLab · Jul 13, 2020 · b3bdf61 · b3bdf61
2 parents 375d9ef + 97eea42
commit b3bdf61
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/robustness/train.py b/robustness/train.py
@@ -118,7 +118,7 @@ def lr_func(ep):
             for i in range(steps_to_take):
                 schedule.step()
 
-        if 'amp' in checkpoint and checkpoint['amp'] is not None:
+        if 'amp' in checkpoint and checkpoint['amp'] not in [None, 'N/A']:
             amp.load_state_dict(checkpoint['amp'])
 
         # TODO: see if there's a smarter way to do this