Only save and load master params if training with FP16

examples/imagenet/main.py

@@ -149,6 +149,7 @@ def main():
                 args.start_epoch = checkpoint['epoch']
                 best_prec1 = checkpoint['best_prec1']
                 model.load_state_dict(checkpoint['state_dict'])
+                if args.fp16:
                     saved_master_params = checkpoint['master_params']
                     for master, saved in zip(master_params, saved_master_params):
                         master.data.copy_(saved.data) 
@@ -219,14 +220,19 @@ def main():
         if args.local_rank == 0:
             is_best = prec1 > best_prec1
             best_prec1 = max(prec1, best_prec1)
-            save_checkpoint({
+            # Use local scope to avoid dangling references
+            def create_and_save_checkpoint():
+                checkpoint_dict = {
                     'epoch': epoch + 1,
                     'arch': args.arch,
                     'state_dict': model.state_dict(),
                     'best_prec1': best_prec1,
                     'optimizer' : optimizer.state_dict(),
-                'master_params': master_params, 
-            }, is_best)
+                }
+                if args.fp16:
+                    checkpoint_dict['master_params'] = master_params
+                save_checkpoint(checkpoint_dict, is_best)
+            create_and_save_checkpoint()
 
 class data_prefetcher():
     def __init__(self, loader):

examples/imagenet/main_reducer.py

@@ -149,6 +149,7 @@ def main():
                 args.start_epoch = checkpoint['epoch']
                 best_prec1 = checkpoint['best_prec1']
                 model.load_state_dict(checkpoint['state_dict'])
+                if args.fp16:
                     saved_master_params = checkpoint['master_params']
                     for master, saved in zip(master_params, saved_master_params):
                         master.data.copy_(saved.data) 
@@ -219,14 +220,19 @@ def main():
         if args.local_rank == 0:
             is_best = prec1 > best_prec1
             best_prec1 = max(prec1, best_prec1)
-            save_checkpoint({
+            # Use local scope to avoid dangling references
+            def create_and_save_checkpoint():
+                checkpoint_dict = {
                     'epoch': epoch + 1,
                     'arch': args.arch,
                     'state_dict': model.state_dict(),
                     'best_prec1': best_prec1,
                     'optimizer' : optimizer.state_dict(),
-                'master_params': master_params, 
-            }, is_best)
+                }
+                if args.fp16:
+                    checkpoint_dict['master_params'] = master_params
+                save_checkpoint(checkpoint_dict, is_best)
+            create_and_save_checkpoint()
 
 class data_prefetcher():
     def __init__(self, loader):