fix FutureWarning with torch.load (#27)

lucidrains · Aug 19, 2024 · 884ec97 · 884ec97
1 parent a6ac238
commit 884ec97
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/ema_pytorch/post_hoc_ema.py b/ema_pytorch/post_hoc_ema.py
@@ -389,7 +389,7 @@ def synthesize_ema_model(
 
             # load checkpoint into a temporary ema model
 
-            ckpt_state_dict = torch.load(str(checkpoint))
+            ckpt_state_dict = torch.load(str(checkpoint), weights_only=True)
             tmp_ema_model.load_state_dict(ckpt_state_dict)
 
             # add weighted checkpoint to synthesized