fix evaluate.py load model

ruiliu-ai · Dec 10, 2021 · e7e461f · e7e461f
1 parent 3f47e17
commit e7e461f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/evaluate.py b/evaluate.py
@@ -257,7 +257,7 @@ def main_worker():
     model = net.InpaintGenerator().to(device)
     model_path = args.ckpt
     data = torch.load(args.ckpt, map_location=device)
-    model.load_state_dict(data['netG'])
+    model.load_state_dict(data)
     print('loading from: {}'.format(args.ckpt))
     model.eval()