Solved by JuanFMontesinos in post #2 
              
              
                Hi,  nn.linear work with an arbitrary amount of dimensions, namely you can pass whatever tensor of size BATCH,,dim yo obtain BATCH,,new_dim. 
Never do for loops in pytorch as it is equivalent to generate Siamese modules. It duplicates the computational graph as many times as you call the module. 
If…