Hi @ptrblck
Does lazy definition work?
class Model(nn.Module):
    def __init__(self, num_classes):
        super().__init__()
        self.conv = nn.Conv2d(3, 6, 3, 1, 1)
        self.linear = None 
        self.num_classes = num_classes

    def forward(self, x):
        x = self.conv(x)
        x = x.view(x.size(0), -1)
        if self.linear is None:
            self.linear = nn.Linear(x.size(1), self.num_classes)
        x = self.linear(x)
        return x