[Fix] dict_keys error in meta_rcnn_roi_head (#6)

open-mmlab · Nov 24, 2021 · ded7c35 · ded7c35
1 parent 8ec0695
commit ded7c35
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/mmfewshot/detection/models/roi_heads/meta_rcnn_roi_head.py b/mmfewshot/detection/models/roi_heads/meta_rcnn_roi_head.py
@@ -343,12 +343,12 @@ def simple_test_bboxes(
         cls_scores_dict[num_classes] /= len(support_feats_dict.keys())
         cls_scores = [
             cls_scores_dict[i] if i in cls_scores_dict.keys() else
-            torch.zeros_like(cls_scores_dict[cls_scores_dict.keys()[0]])
+            torch.zeros_like(cls_scores_dict[list(cls_scores_dict.keys())[0]])
             for i in range(num_classes + 1)
         ]
         bbox_preds = [
             bbox_preds_dict[i] if i in bbox_preds_dict.keys() else
-            torch.zeros_like(bbox_preds_dict[bbox_preds_dict.keys()[0]])
+            torch.zeros_like(bbox_preds_dict[list(bbox_preds_dict.keys())[0]])
             for i in range(num_classes)
         ]
         cls_score = torch.cat(cls_scores, dim=1)