fix gardcam error using in mmdet (#779)

open-mmlab · May 17, 2023 · 1e5372e · 1e5372e
1 parent a3b6ae1
commit 1e5372e
Showing 1 changed file with 8 additions and 3 deletions.
diff --git a/mmyolo/utils/boxam_utils.py b/mmyolo/utils/boxam_utils.py
@@ -202,8 +202,10 @@ def __call__(self, *args, **kwargs):
         if self.is_need_loss:
             # Maybe this is a direction that can be optimized
             # self.detector.init_weights()
-
-            self.detector.bbox_head.head_module.training = True
+            if hasattr(self.detector.bbox_head, 'head_module'):
+                self.detector.bbox_head.head_module.training = True
+            else:
+                self.detector.bbox_head.training = True
             if hasattr(self.detector.bbox_head, 'featmap_sizes'):
                 # Prevent the model algorithm error when calculating loss
                 self.detector.bbox_head.featmap_sizes = None
@@ -219,7 +221,10 @@ def __call__(self, *args, **kwargs):
 
             return [loss]
         else:
-            self.detector.bbox_head.head_module.training = False
+            if hasattr(self.detector.bbox_head, 'head_module'):
+                self.detector.bbox_head.head_module.training = False
+            else:
+                self.detector.bbox_head.training = False
             with torch.no_grad():
                 results = self.detector.test_step(self.input_data)
                 return results