I don’t know if you are flattening the activation before passing it to the first linear layer, but you are hitting the same issue other users were running into in this thread: the feature dimension of the activation doesn’t fit the in_features of the first linear layer.
If you are already flattening the activation, change in_features to 1024, if not flatten the activation and see if it’s working.
Also, it’s better to post code snippets by wrapping them into three backticks ``` as it makes debugging easier.