I have a tensor x of shape [batch_size, 2]. I want to split it in shapes [batch_size, 1] and [batch_size, 1] then use one as a classification head and one as a regression head.
Is there a way to do this or will I have to use different tensors?
Use the tensors x[:, 0] and x[:, 1].
I did that but I am asking that if I pass these to BCE loss and MSE loss respectively, will the grad be stored correctly so that optimisation can happen properly?
Yes ! Say you have loss_1 and loss_2 (computed for example with x[:, 0] and x[:, 1]).
Then if you call .backward() on a combination of the two losses, for example
loss_sum = 0.5 * loss_1 + 0.5 * loss_2
loss_sum.backward()
optimization can happen properly.