Yes, it works now! And I solved the second problem checking your nice answer here (x.view(x.size(0), -1) indeed).
Now, for the next module I need to stop the current gradient since I want to use local backprop only, is the correct approach to just use detach?