Skip to content

Commit aa2a2c1

Browse files
committed
add pytorch < 1.13 compatibility
1 parent 0b5f619 commit aa2a2c1

File tree

1 file changed

+5
-1
lines changed

1 file changed

+5
-1
lines changed

bioimageio/core/backends/pytorch_backend.py

+5-1
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
from typing_extensions import assert_never
1313

1414
from bioimageio.spec._internal.type_guards import is_list, is_ndarray, is_tuple
15+
from bioimageio.spec._internal.version_type import Version
1516
from bioimageio.spec.common import ZipPath
1617
from bioimageio.spec.model import AnyModelDescr, v0_4, v0_5
1718
from bioimageio.spec.utils import download
@@ -143,7 +144,10 @@ def load_torch_state_dict(
143144
model = model.to(devices[0])
144145
with path.open("rb") as f:
145146
assert not isinstance(f, TextIOWrapper)
146-
state = torch.load(f, map_location=devices[0], weights_only=True)
147+
if Version(str(torch.__version__)) < Version("1.13"):
148+
state = torch.load(f, map_location=devices[0])
149+
else:
150+
state = torch.load(f, map_location=devices[0], weights_only=True)
147151

148152
incompatible = model.load_state_dict(state)
149153
if (

0 commit comments

Comments
 (0)