From the example you gave, I gather that we need a tensor with length equals to our training dataset and with a values equal to the class weight. as if I have this dataset [0,0,0,0,1,1] then the weight value for each sample would be [0.25,0.25,0.25,0.25,0.5,0.5].
Is there a simple way to get this weight value for each sample knowing the ammount of samples of each class? I could not reproduce your sample code for my 9 class example.