File tree 1 file changed +5
-1
lines changed
1 file changed +5
-1
lines changed Original file line number Diff line number Diff line change 12
12
from typing_extensions import assert_never
13
13
14
14
from bioimageio .spec ._internal .type_guards import is_list , is_ndarray , is_tuple
15
+ from bioimageio .spec ._internal .version_type import Version
15
16
from bioimageio .spec .common import ZipPath
16
17
from bioimageio .spec .model import AnyModelDescr , v0_4 , v0_5
17
18
from bioimageio .spec .utils import download
@@ -143,7 +144,10 @@ def load_torch_state_dict(
143
144
model = model .to (devices [0 ])
144
145
with path .open ("rb" ) as f :
145
146
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 )
147
151
148
152
incompatible = model .load_state_dict (state )
149
153
if (
You can’t perform that action at this time.
0 commit comments