PyTorch 2 · Autograd and models
Define a model
Subclass nn.Module, create layers in __init__ and define forward.
class Net(nn.Module):
def __init__(self):
super().__init__()
self.layers = nn.Sequential(
nn.Linear(784, 128),
nn.ReLU(),
nn.Linear(128, 10),
)
def forward(self, x):
return self.layers(x)
Official PyTorch 2 documentation