Merge pull request #46 from johahi/main

small fix to disable future warning with torch.load
lucidrains · Oct 9, 2024 · 31f61d7 · 31f61d7
2 parents d4ae222 + a94ec15
commit 31f61d7
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/enformer_pytorch/modeling_enformer.py b/enformer_pytorch/modeling_enformer.py
@@ -26,7 +26,7 @@
 # solution came from @johahi
 
 DIR = Path(__file__).parents[0]
-TF_GAMMAS = torch.load(str(DIR / "precomputed"/ "tf_gammas.pt"))
+TF_GAMMAS = torch.load(str(DIR / "precomputed"/ "tf_gammas.pt"), weights_only=True)
 
 # helpers