diff --git a/intermediate_source/torchvision_tutorial.py b/intermediate_source/torchvision_tutorial.py
index a19d393043..b235df7578 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])