import torch
import torch.nn as nn

# Random data: 100 examples, each with 3 features
X = torch.randn(100, 3)
# Random "labels" (what the network should learn to predict)
y = torch.randn(100, 1)

# Define a simple single-layer neural network
model = nn.Linear(3, 1)

# Mean squared error loss function
loss_fn = nn.MSELoss()

# Stochastic Gradient Descent optimizer
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

# Training loop
for epoch in range(100):
    # Forward pass: compute prediction
    y_pred = model(X)
    # Compute loss
    loss = loss_fn(y_pred, y)
    # Zero gradients, backward pass, and update weights
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    if (epoch+1) % 20 == 0:
        print(f'Epoch {epoch+1}, Loss: {loss.item():.4f}')

# After training, model(X) gives outputs closer to y