fix: load checkpoint on right device

atong01 · May 20, 2024 · 333d73f · 333d73f
1 parent d0b0da2
commit 333d73f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/examples/images/cifar10/compute_fid.py b/examples/images/cifar10/compute_fid.py
@@ -51,7 +51,7 @@
 # Load the model
 PATH = f"{FLAGS.input_dir}/{FLAGS.model}/{FLAGS.model}_cifar10_weights_step_{FLAGS.step}.pt"
 print("path: ", PATH)
-checkpoint = torch.load(PATH)
+checkpoint = torch.load(PATH, map_location=device)
 state_dict = checkpoint["ema_model"]
 try:
     new_net.load_state_dict(state_dict)