Fix loading for .pth checkpoints

Xilinx · Oct 15, 2019 · 12cf2f1 · 12cf2f1
1 parent a94daa8
commit 12cf2f1
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/examples/imagenet_val.py b/examples/imagenet_val.py
@@ -51,7 +51,7 @@ def main():
  # Map model to be loaded to specified single gpu.
  loc = 'cuda:{}'.format(args.gpu)
  checkpoint = torch.load(args.resume, map_location=loc)
- model.load_state_dict(checkpoint['state_dict'], strict=False)
+ model.load_state_dict(checkpoint, strict=False)
 
  valdir = os.path.join(args.imagenet_dir, 'val')
  normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])