Solved by ptrblck in post #2 
              
              
                The linear layer is pretty straightforward, as the output activation will have the shape [batch_size, *, out_features]. 
The activation of the conv layer is a bit trickier, as the shape depend on the spatial size of your input: 
x = torch.randn(2, 10)
out = l(x)
print(out.nelement())  # batch_size *…