From 46fe1b110e862b222848f3cf011a5c226a408122 Mon Sep 17 00:00:00 2001 From: shoushinya Date: Sat, 13 Jun 2026 19:57:06 +0800 Subject: [PATCH] fix: use torch.as_tensor for Mask-RCNN labels to support multiple classes --- intermediate_source/torchvision_tutorial.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/intermediate_source/torchvision_tutorial.py b/intermediate_source/torchvision_tutorial.py index a19d3930436..b235df7578e 100644 --- a/intermediate_source/torchvision_tutorial.py +++ b/intermediate_source/torchvision_tutorial.py @@ -173,7 +173,7 @@ def __getitem__(self, idx): boxes = masks_to_boxes(masks) # there is only one class - labels = torch.ones((num_objs,), dtype=torch.int64) + labels = torch.as_tensor(obj_ids, dtype=torch.int64) image_id = idx area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])