Merge pull request #5 from boostcampaitech4lv23cv3/train-csc

Modify train.py
boostcampaitech4lv23cv3 · Nov 28, 2022 · c80e0d5 · c80e0d5
2 parents e75ad1a + 4a879fe
commit c80e0d5
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/mmdetection/tools/train.py b/mmdetection/tools/train.py
@@ -205,7 +205,8 @@ def main():
     logger = get_root_logger(log_file=log_file, log_level=cfg.log_level)
 
     # set num_classes
-    cfg.model.roi_head.bbox_head.num_classes = len(cfg.classes)
+    for bbox_head in cfg.model.roi_head.bbox_head:
+        bbox_head.num_classes = len(cfg.classes)
 
     # init the meta dict to record some important information such as
     # environment info and seed, which will be logged