MulitVariable Linear Regression using Pytorch

[PrathPrath]

Hey all, I just started with my NN journey using pytorch and was experimenting with the sklearn.datasets.load_diabetes dataset and was trying design a multivariable linear regression model but unfortunately I see to be stuck.
Can anyone help me here?

`data = load_diabetes()
X = data.data
y = data.target
X = StandardScaler().fit_transform(X)

X = torch.tensor(X,dtype=torch.float32)
Y = torch.tensor(Y,dtype=torch.float32)

X_train, X_test, Y_train, Y_test = train_test_split(X,Y,test_size=0.2)

class MultiLinearRegression(nn.Module):
def init(self):
super().init()
self.Linear = nn.Linear(10,1)

def forward(self,x:torch.tensor) -> torch.tensor:
return self.Linear(x)

torch.manual_seed(42)
model = MultiLinearRegression()

loss_fn = nn.MSELoss()
optimizer = torch.optim.SGD(params=model.parameters(),lr=0.2)

torch.manual_seed(42)

epochs = 1000
epoch_count =[]
train_loss_list =[]
test_loss_list =[]

for epoch in range(epochs):
model.train()

y_pred_train = model(X_train)
loss_train = loss_fn(y_pred_train,Y_train)

optimizer.zero_grad()

loss_train.backward()

optimizer.step()

model.eval()

with torch.inference_mode():
y_pred_test = model(X_test)
loss_test = loss_fn(y_pred_test,Y_test)

if epoch%10==0:
  # append all required data for later evaluation
  epoch_count.append(epoch)
  train_loss_list.append(loss_train)
  test_loss_list.append(loss_test)

print(f"Epoch : {epoch} | Loss :{loss_train} | Test_loss : {loss_test}")
print(model.state_dict())

plot the loss curve

plt.plot(epoch_count,test_loss_list,color='g',label="Test Loss Cruve")
plt.plot(epoch_count,torch.tensor(train_loss_list).numpy(),color='b',label="Train Cruve")
plt.legend()
plt.show()

model.eval()
with torch.inference_mode():
y_pred = model(X_test)

`

I have played around with the lr without any luck

`y_pred[:5],Y_test[:5]

(tensor([[148.9813],
[148.9791],
[148.9810],
[148.9803],
[148.9744]]),
tensor([[[259.]],

     [[ 52.]],

     [[198.]],

     [[ 72.]],

     [[248.]]]))`