Update model/backbones.py

model/backbones.py

@@ -1,14 +1,19 @@
 import torch.nn as nn
-import torchvision.models as models
+from torchvision.models import mobilenet_v2, MobileNet_V2_Weights
 
+class MobileNetV2Backbone(nn.Module):
+    def __init__(self):
+        super(MobileNetV2Backbone, self).__init__()
+        # Используем актуальный способ загрузки предобученной модели
+        self.model = mobilenet_v2(weights=MobileNet_V2_Weights.DEFAULT).features
+        self.enc_channels = [24, 32, 96, 320]  # Каналы с разных уровней feature extractor
 
-def _build_mobilenetv2_backbone():
-    backbone = models.mobilenet_v2(pretrained=True).features
-    backbone.enc_channels = [24, 32, 96, 160, 320]
-    return backbone
+    def forward(self, x):
+        return self.model(x)
 
 
+# Добавим SUPPORTED_BACKBONES прямо сюда (или перенеси в backbones/__init__.py, если нужно)
 SUPPORTED_BACKBONES = {
-    'mobilenetv2': lambda in_channels: _build_mobilenetv2_backbone()
-}
+    'mobilenetv2': lambda _=None: MobileNetV2Backbone()
+}  # _=None позволяет игнорировать входной параметр