from easytorch.layer import Linear, Tanh, Sequential
from easytorch.optim import SGD
import easytorch.functional as F

# Create a model, optimizer, loss function
model = Sequential(
    Linear(1, 5),
    Tanh(),
    Linear(5, 1)
)
opt = SGD(model.parameters(), lr=3e-4)
loss_fn = F.mse_loss

# train the model
for epoch in range(epochs):
    pred = model(x)
    loss = loss_fn(pred, y)
    opt.zero_grad()
    loss.backward()
    opt.step()