class LinearRegressionModel(nn.Module): def __init__(self): super().__init__() self.weights = nn.Parameter(torch.randn(1, requires_grad=True, dtype=torch.float)) self.bias = nn.Parameter(torch.randn(1, requires_grad=True, dtype = torch.float)) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.weights * x + self.bias model = LinearRegressionModel()
epochs = 300 for epoch in range(epochs): model.train() #set up the mode. Train mode in Pytorch sets all parameter to requite grad that require the gradients y_pred = model(X_train) #forward pass loss = loss_fn(y_pred, y_train) # calculate loss optimizer.zero_grad()# Assign 0 grad so gradients don't get accumulated at each step loss.backward() # perform back propagation optimizer.step() # Update the parameters according to the gradients. model.eval() #turn off settings not needed for evaluation with torch.inference_mode(): test_pred = model(X_test) test_loss = loss_fn(test_pred,y_test)