Thank you all! I think it works better now
class Custom_MSELoss(torch.nn.Module):
    def __init__(self):
        super(Custom_MSELoss, self).__init__()

    def forward(self, output, x_train, x_expected):
        x_train = torch.sum(output) * 0 + x_train  # <--!
        MSE = [(x_train[0][i] - x_expected[0][i])**2 for i in range(len(x_expected[0]))]
        MSE = sum(MSE) / len(MSE)
        return MSE

criterion = Custom_MSELoss()

Looks fine to me. Could any operation here potentially detach the history?