I am new to pytorch and have a question about grad calculating in terms about multi-task training.

For example, I need to train Classification with data label == 1 or 0; At the same time, I need to train Regression with data label == 1. The data format for label == 1 is [batch, class 1, x1, x2, y1. y2] and for label == 0 is [batch, class 0, 0].

My ouput tensor for classification is out_cls [batch, class] (like [100, 2]) and for regression is out_reg [batch, pred_x1, pred_x2, pred_y1, pred_y2] (like [100 ,4]).

I am expecting the final loss = loss_cls + loss_reg, then loss.bardword().

What am I doing now
Get the valid index for classification and regression. Get the valid tensor by out_cls = out_cls[lvalid_index], out_reg = out_reg[valid_index]. Then use these new tensor to calculate the loss, like loss_cls = nn.CrossEntrophy(out_cls, target_cls). loss_reg = nn.MSE(out_reg, target_reg).

Is this a good practice for multi-task training ? Am I successfully avoiding calculate grad for regression with data label == 0 ？