replacing
labels = torch.ones((num_objs,), dtype=torch.int64)

with the following
labels = torch.as_tensor(obj_ids, dtype=torch.int64)
will do.
In fact, the replacement even works for two-class instance segmentation.