You can flatten the activations via x = x.view(x.size(0), -1) or by using nn.Flatten.
However, I don’t think flattening the activations alone will solve the issue, since the model output should have the shape [batch_size, nb_classes] for a multi-class classification as shown in the tutorial.
Thus I would still recommend to use e.g. an nn.Linear module to create the desired outputs.