Update yolov5/models/experimental.py

yolov5/models/experimental.py

@@ -3,16 +3,15 @@
 Experimental modules
 """
 import math
-import models.yolo
 import numpy as np
 import torch
 import torch.nn as nn
-import torch
 
-from models.common import Conv  # as the error previously asked for
-from models.yolo import DetectionModel
+import models.yolo  # ensures module resolution
+from models.yolo import DetectionModel  # for allowlisting
 from torch.nn.modules.container import Sequential
-from models.common import Conv  # <--- new: the class in the latest error
+from models.common import Conv  # YOLO-specific common conv block
+
 from utils.downloads import attempt_download
 
 
@@ -82,10 +81,22 @@ def attempt_load(weights, device=None, inplace=True, fuse=True):
     model = Ensemble()
 
     for w in weights if isinstance(weights, list) else [weights]:
-        with torch.serialization.safe_globals([DetectionModel, Sequential, Conv, torch.nn.modules.conv.Conv2d]):
-            ckpt = torch.load(attempt_download(w), map_location='cpu')  # safe load
+        # try safe allowlisted load first
+        try:
+            with torch.serialization.safe_globals([
+                DetectionModel,
+                Sequential,
+                Conv,
+                torch.nn.modules.conv.Conv2d,
+                torch.nn.modules.batchnorm.BatchNorm2d
+            ]):
+                ckpt = torch.load(attempt_download(w), map_location='cpu')  # safe load
+        except Exception:
+            # fallback to full load if checkpoint is trusted
+            ckpt = torch.load(attempt_download(w), map_location='cpu', weights_only=False)
+
         ckpt = (ckpt.get('ema') or ckpt['model']).to(device).float()  # FP32 model
-    
+
         # Model compatibility updates
         if not hasattr(ckpt, 'stride'):
             ckpt.stride = torch.tensor([32.])