Thanks for the clarification!
In that case you would have to deal with some edge cases:

if the current batch does not contain samples from a specific class, you should skip it. Otherwise you will get a all-zero mask and the code will raise an exception. Maybe you could use torch.unqiue to get all current class labels and loop over these indices instead of the loop for all classes
if the current batch only contains a single sample of a class, your nn.BatchNorm1d layer(s) will throw an exception, since they cannot calculate the batch stats for a single sample (and 1 sequence length).

However, if you somehow make sure that your batches are balanced, your code works fine:
model = my_model()
print(model)
target = torch.randn(60,10)
loss =nn.MSELoss()
input = torch.randn(60,10)
labels = torch.arange(10).view(10, 1).repeat(1, 6).view(-1)
zn = torch.randn(60,10)
predict = model(zn, labels)
output = torch.cat(predict)
criteria = loss(output,target)