Splitting Tensor for different loss functions

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

optimization can happen properly.