I have a regression model with multiple outputs, so I have a custom mse function
def custom_mse(predicted, target):
    total_mse = 0
    for i in range(target.shape[1]):
        total_mse+=nn.MSELoss()(predicted[i], target[i])
    return total_mse

the purpose is to add up all the mse for each individual output, and return the total as the loss. This works well enough, but is there a way to get rid of the for loop and vectorize the total_mse calculation?