Thank you very much for your quick reply! I’ve tried adding this code, filling in the blanks as follows:
class ConvNet(nn.Module):

  def __init__(self, n_classes):
    super.__init__()
    self.featurizer = nn.Sequential(
              nn.Conv2d(3, 32, kernel_size=3, stride=2, padding=1),
              nn.ReLU(),
              nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1),
              nn.ReLU(),
              nn.Conv2d(64, 21, kernel_size=3, stride=2, padding=1),
              nn.ReLU(),
              nn.MaxPool2d(1)) # what you  have in your model variable 
    self.classifier = nn.Sequential(
              nn.Linear(21*3*3, 120),
              nn.Sigmoid(),
              nn.Linear(120, 84),
              nn.Sigmoid(),
              nn.Linear(84, 21)
    ) # 2-3 linear layers, first should have in_features = n_kernels * h_kernel * w_kernel from last conv layer, last should have out_feautres=n_classes )

  def forward(self, x):
    # x (b, c, h, w)
    x = self.featurizer(x) # x (b, n_kernels, h_kernel, w_kernel)
    x = torch.flatten(x, start_dim=1) # x (b, n_kernels * h_kernel, w_kernel)
    x = self.classifier(x) # x (b, n_classes)
    x = torch.nn.f.softmax(x) 

I’m afraid to speak too soon but the permute seems to no longer be necessary now! I do, however, get the following error still:
RuntimeError: mat1 and mat2 shapes cannot be multiplied (10x21504 and 189x120)
The 189 x 120 are clearly the parameters I set in the first linear layer, and I assume this error refers to a mismatch between these parameters and the data it actually receives. Which does however make me quite confused as to where the 10 x 21504 comes from.
Is something in my convolutional layers shaping the data in these dimensions?
Edit: Or the “flatten” layer, I just realized. But there doesn’t really seem to be a way to control its output dimensions other than start_dim and end_dim.