support new version of torchvision (median-research-group#62)

biubiuer · Nov 22, 2023 · d300730 · d300730
1 parent 5f81dfa
commit d300730
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/LibMTL/model/resnet.py b/LibMTL/model/resnet.py
@@ -1,6 +1,9 @@
 import torch
 import torch.nn as nn
-from torchvision.models.utils import load_state_dict_from_url
+try:
+    from torchvision.models.utils import load_state_dict_from_url
+except:
+    from torch.hub import load_state_dict_from_url
 
 
 # __all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101',