Update models/experimental.py

models/experimental.py

@@ -251,11 +251,12 @@ def attempt_load(weights, map_location=None):
         attempt_download(w)
         ckpt = torch.load(w, map_location=map_location)  # load
 
-        # 处理没有 'ema' 键的情况
-        state_dict = ckpt['model'] if 'model' in ckpt else ckpt
-
-        # 将 state_dict 作为模型参数
-        model.append(state_dict.float().eval())  # 不使用 .fuse()
+        if isinstance(ckpt, dict):
+            state_dict = ckpt['model'] if 'model' in ckpt else ckpt
+            model.append(state_dict.float().eval())  # 不使用 .fuse()
+        else:
+            # 如果 ckpt 是 GANLearner 实例，直接添加到模型中
+            model.append(ckpt)
 
     # Compatibility updates
     for m in model.modules():
@@ -272,4 +273,4 @@ def attempt_load(weights, map_location=None):
         print('Ensemble created with %s\n' % weights)
         for k in ['names', 'stride']:
             setattr(model, k, getattr(model[-1], k))
-        return model  # return ensemble
+        return model  # return ensemble