PyTorch: simple example

simple-nn-pytorch.py

import torch
import torch.nn as nn
import torch.nn.functional as fun

data_size = 10

input_size = 28 * 28
hidden1_output = 200
output_size = 1

data = torch.randn(data_size, input_size)
target = torch.randn(data_size, output_size)

model = nn.Sequential(
    nn.Linear(input_size, hidden1_output),
    nn.ReLU(),
    nn.Linear(hidden1_output, output_size)
)

opt = torch.optim.SGD(model.parameters(), lr=1e-3)

for step in range(100):
    target_ = model(data)
    loss = fun.mse_loss(target_, target)
    loss.backward()
    opt.step()
    opt.zero_grad()