Yes I did! I think the index attribute expects a tensor of all the classes for example if I have 10 classes, then it expects a tensor([0,1,2,3,4,5,6,7,8,9]). So I declared a tensor in this manner and it worked for me. Let me know if it works for you too.