This says that a linear layer gets the wrong number of features (and you’d know from the backtrace which). I would propose to print all shapes of interest in the forward to check on it.
Quite likely this is the average pool not doing what you think it does, maybe AdaptiveAveragePool (or just calling “features[0].mean([2, 3])”) is a better fit for your purpose.
Best regards
Thomas