I have a classification problem where each sample output consists of four values each of which is one of the 6 classes, i.e y has shape (batchsize, 4). So far, I have created a toy model which outputs a logits tensor out having the shape (batchsize, 6(=num_classes), 4). I am calculating the loss for each of the target value and summing up, followed by backprop:
out = model(in)
out.shape #(batchsize, 6, 4)

out1 = out[:,:,0].squeeze(-1)
out1.shape #(batchsize, 6)
out2 = out[:,:,1].squeeze(-1)
out3 = out[:,:,2].squeeze(-1)
out4 = out[:,:,3].squeeze(-1)
y1 = y[:,0].squeeze(-1)
y1.shape #(batchsize)
y2 = y[:,1].squeeze(-1)
y3 = y[:,2].squeeze(-1)
y4 = y[:,3].squeeze(-1)

total_loss = criterion(out1, y1) + criterion(out2, y2) + criterion(out3, y3) + criterion(out4, y4)

optimizer.zero_grad()
total_loss.backward()
optimizer.step()

I want to know if this is the correct workaround for multivariate classification