support load_state_dict_from_url in later pytorch versions

segmentation/model/backbone/hrnet.py

@@ -25,7 +25,11 @@
 import logging
 import torch.nn as nn
 import torch.nn.functional as F
-from torchvision.models.utils import load_state_dict_from_url
+
+try:
+    from torch.hub import load_state_dict_from_url
+except:
+    from torchvision.models.utils import load_state_dict_from_url
 
 logger = logging.getLogger('hrnet_backbone')
 

segmentation/model/backbone/mnasnet.py

@@ -8,7 +8,11 @@
 
 import torch
 import torch.nn as nn
-from torchvision.models.utils import load_state_dict_from_url
+
+try:
+    from torch.hub import load_state_dict_from_url
+except:
+    from torchvision.models.utils import load_state_dict_from_url
 
 __all__ = ['MNASNet', 'mnasnet0_5', 'mnasnet0_75', 'mnasnet1_0', 'mnasnet1_3']
 

segmentation/model/backbone/mobilenet.py

@@ -4,8 +4,11 @@
 # ------------------------------------------------------------------------------
 
 from torch import nn
-from torchvision.models.utils import load_state_dict_from_url
 
+try:
+    from torch.hub import load_state_dict_from_url
+except:
+    from torchvision.models.utils import load_state_dict_from_url
 
 __all__ = ['MobileNetV2', 'mobilenet_v2']
 

segmentation/model/backbone/resnet.py

@@ -4,8 +4,11 @@
 # ------------------------------------------------------------------------------
 
 import torch.nn as nn
-from torchvision.models.utils import load_state_dict_from_url
 
+try:
+    from torch.hub import load_state_dict_from_url
+except:
+    from torchvision.models.utils import load_state_dict_from_url
 
 __all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101',
            'resnet152', 'resnext50_32x4d', 'resnext101_32x8d',

segmentation/model/backbone/xception.py

@@ -6,7 +6,11 @@
 from collections import OrderedDict
 
 import torch.nn as nn
-from torchvision.models.utils import load_state_dict_from_url
+
+try:
+    from torch.hub import load_state_dict_from_url
+except:
+    from torchvision.models.utils import load_state_dict_from_url
 
 __all__ = ['Xception65', 'xception65']