Multiple Linear Regression results in Nan even when reducing LR and increasing epochs

Hello, I am new to machine learning and pytorch so I am starting with linear regression.

This is my code:

class LinearRegressionModel(torch.nn.Module):
    
    def __init__(self):
        super(LinearRegressionModel, self).__init__()
        self.linear = torch.nn.Linear(4, 1)  # four input and one output
    
    def forward(self, x):
        y_pred = self.linear(x)
        return y_pred
    
our_model = LinearRegressionModel()

#changed size_average = False to reduction = 'sum' because of depracated value
criterion = torch.nn.MSELoss(reduction = 'sum')
optimizer = torch.optim.SGD(our_model.parameters(), lr = 0.01)


#running into issue with data being either double or float
#converted all data into float
for epoch in range(500):
 
    pred_y = our_model(X)
    
    loss = criterion(pred_y, ynew)
    
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    print('epoch {}, loss {}'.format(epoch, loss.item()))

I keep getting returned Nan even when the LR is reduced to increased or when the number of epochs is increased.

You could check the .grad attributes of the parameters and make see if they are blowing up.
If so, reduce the learning rate further and use a mean reduction in the criterion.