If you don’t want to change the backbone, and opt for using resnet50 model described in the tutorial, all should go well if you use labels = torch.as_tensor(obj_ids, dtype=torch.int64). Unless there is something incorrect in your dataset class. Hence, I suggest that you revise your dataset class, best way to do this is by invoking it separately (a simple program with a main function) and check that everything is correct; better off in the debug mode (via Spyder or PyChram) to trace everything in a step-by-step bases. You might as well go over the rest of your code to make sure everything is correct. I have tried it with 60 different classes and it is working.