PyTorch Linear Regression with 2 weights and a bias

[Prithivi-002]

I tried to create a Linear Regression model with 2 independent variable and one dependent variable like Y = xa_1 + za_2 + c... Can anyone tell me what did I do wrong ?

Code:
Dataset:
weight_1 = 0.2
weight_2 = 0.4
bias = 0.5

start = 0
end = 10
step = 0.1
X = torch.arange(start, end, step)

startz = 0
endz = 100
stepz = 1
Z = torch.arange(startz, endz, stepz)

Y = X * weight_1 + Z * weight_2 + bias

X[:10], Z[:10], Y[:10], len(X), len(Z), len(Y)

Training Testing Split
train_split = int(0.8 * len(X))
X_train, Z_train, Y_train = X[:train_split], Z[:train_split], Y[:train_split]
X_test, Z_test, Y_test = X[train_split:],Z[train_split:], Y[train_split:]

x = torch.stack([X_train, Z_train], dim = 1)
y = torch.stack([Y_train], dim = 1)
x_test = torch.stack([X_test, Z_test], dim = 1)
y_test = torch.stack([Y_test], dim = 1)
x = x.view(-1, 2)
y = y.view(-1, 1)
x_test = x_test.view(-1, 2)
y_test = y_test.view(-1, 1)
x.shape, y.shape, x_test.shape, y_test.shape

Pytorch Model

Creating PyTorch Model

class LinerRegression(nn.Module):
def init(self):
super().init()
self.linear_layer = nn.Linear(in_features = 2, out_features = 1)

def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.linear_layer(x)

model Initialization

torch.manual_seed(42)
model = LinerRegression()

model.state_dict()

Loss Function and optimizer

loss_fn = nn.SmoothL1Loss()
optimizer = torch.optim.SGD(params = model.parameters(), lr = 0.001)

Trainig and Testing Loop

torch.manual_seed(42)

epochs = 100

for epoch in range(epochs):
model.train()
y_pred = model(x)

loss = loss_fn(y_pred, y)
optimizer.zero_grad()
loss.backward()
optimizer.step()

testing

model.eval()
with torch.inference_mode():
test_pred = model(x_test)
test_loss = loss_fn(test_pred, y_test)
test_loss_values.append(test_loss)

Append epoch_count and train_loss_values

epoch_count.append(epoch)
train_loss_values.append(loss)

if epoch % 10 == 0:
print(f"Epoch: {epoch} | Loss: {loss} | Test Loss: {test_loss}")

I like to create a linear regression model with 2 weights and a bias.. tell me where did I go wrong..

[LuluW8071]

@Prithivi-002

# Insert empty list above train loop
train_loss_values = []
test_loss_values = []
epoch_count = []

# Set lower learning rate
optimizer = torch.optim.SGD(params = model.parameters(), lr = 0.0001)