From 884ec97d6cd11428509e2b958ae204f42ed637d6 Mon Sep 17 00:00:00 2001 From: Dolen <87627379+HiDolen@users.noreply.github.com> Date: Mon, 19 Aug 2024 09:55:22 +0800 Subject: [PATCH] fix FutureWarning with `torch.load` (#27) --- ema_pytorch/post_hoc_ema.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ema_pytorch/post_hoc_ema.py b/ema_pytorch/post_hoc_ema.py index 7994078..8e20442 100644 --- 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