Thanks a lot @ptrblck !
You are right, the error came from within the Encoder… The solution you proposed works, and it seems like the following might work as well, for variable c_hid sizes:
nn.Linear(2*644*c_hid, latent_dim)