Update utils.py

utils.py

@@ -64,8 +64,8 @@ class Dataset(object):
         
         boxes = extract_bboxes(masks)
         boxes = torch.as_tensor(boxes, dtype=torch.float32)
-        # there is only one class
-        labels = torch.ones((num_objs,), dtype=torch.int64)
+        # two_classes
+        labels = torch.as_tensor(class_ids, dtype=torch.int64)
         masks = torch.as_tensor(masks, dtype=torch.uint8)
 
         image_id = torch.tensor([idx])