Specify num_classes = 2 for my custom dataset

facebookresearch#108 (comment)

models/detr.py

@@ -315,6 +315,10 @@ def build(args):
         # for panoptic, we just add a num_classes that is large enough to hold
         # max_obj_id + 1, but the exact value doesn't really matter
         num_classes = 250
+    if args.dataset_file == 'custom':
+        # "You should always use num_classes = max_id + 1 where max_id is the highest class ID that you have in your dataset."
+        # Reference: https://github.com/facebookresearch/detr/issues/108#issuecomment-650269223
+        num_classes = 2
     device = torch.device(args.device)
 
     backbone = build_backbone(args)