nn.Module

nn.Module is the base class of every PyTorch model. You declare the layers in __init__ and the computation in forward. Parameters are registered automatically, which makes it possible to list them, move them to the GPU, and save the model's state.

Example

class Reseau(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(10, 1)
    def forward(self, x): return self.fc(x)

Related concepts

The whole glossary