Solved by ptrblck in post #2 
              
              
                Based on the error message, the activation input has 32 features so you would need to use in_features=32 in the linear layer: 
lin = nn.Linear(32 * 32 * 3, 32)
x = torch.randn(32 * 32 * 3, 32)
out = lin(x)
> RuntimeError: mat1 and mat2 shapes cannot be multiplied (3072x32 and 3072x32)


# works
lin …