I have a problem where I am regressing over a vector of length 8. The first four values of the vector represent one output and the next four represent another output. Since these are two different quantities with different scales, I want to split my output into two, such that the loss is calculated individually for each output quantity and then summed up when backpropagated.
Can I write something like this in the forward function
def forward(self, x):
    x = self.layer1(x) # feature computation
    ...
    x = self.fc(x) #last layer (fully connected)
    y, z = x.split(2)
    return y, z

I don’t know if it will be handled well by the loss function or not. And, does it also require two nn.Linear layers at the end in parallel?