Load with weights_only=True (#68)

dftd3 · Sep 19, 2024 · 8f78d19 · 8f78d19
1 parent d928ba4
commit 8f78d19
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 2 deletions.
diff --git a/src/tad_dftd3/data/radii.py b/src/tad_dftd3/data/radii.py
@@ -26,6 +26,7 @@
 from typing import Optional
 
 import torch
+from tad_mctc._version import __tversion__
 from tad_mctc.data.radii import COV_D3
 
 from ..typing import Tensor
@@ -51,8 +52,12 @@ def _load_vdw_rad_d3(
     Tensor
         VDW radii.
     """
+    kwargs: dict = {"map_location": device}
+    if __tversion__ > (1, 12, 1):  # pragma: no cover
+        kwargs["weights_only"] = True
+
     path = op.join(op.dirname(__file__), "vdw-d3.pt")
-    return torch.load(path).type(dtype).to(device)
+    return torch.load(path, **kwargs).type(dtype=dtype)
 
 
 VDW_D3 = _load_vdw_rad_d3()
diff --git a/src/tad_dftd3/reference.py b/src/tad_dftd3/reference.py
@@ -23,6 +23,7 @@
 from typing import Optional
 
 import torch
+from tad_mctc._version import __tversion__
 
 from .typing import Any, NoReturn, Tensor, get_default_device, get_default_dtype
 
@@ -177,8 +178,12 @@ def _load_c6(
     Tensor
         Reference C6 coefficients.
     """
+    kwargs: dict = {"map_location": device}
+    if __tversion__ > (1, 12, 1):  # pragma: no cover
+        kwargs["weights_only"] = True
+
     path = op.join(op.dirname(__file__), "reference-c6.pt")
-    return torch.load(path).type(dtype).to(device)
+    return torch.load(path, **kwargs).type(dtype=dtype)
 
 
 class Reference: