I forgot to post the (very simple) class I was using to test how the batch computation worked, in case anyone wants to try it here it is:
class linearRegression(nn.Module):

    def __init__(self):
        nn.Module.__init__(self)
        self.params = torch.nn.Parameter(torch.nn.init.normal_(torch.empty(2, 1).float(), mean=0.0, std=0.1))

        self.A = torch.tensor([1, 0], requires_grad=False, dtype=torch.float32).view(1, 2)
        self.B = torch.tensor([0, 1], requires_grad=False, dtype=torch.float32).view(1, 2)

    def forward(self, obs):
        y = (self.A @ self.params) * obs + self.B @ self.params

        return y

It implements the line equation: y = a*x + b, where I want to train a and b given input x and y_ground_truth.
criterion = torch.nn.MSELoss(reduction='mean')

Hope it hepls!