Solved by G.M in post #2 
              
              
                .cuda() only transfers the registered parameters of a model to the gpu, since a is a tensor created locally in forward, u have to transfer it to gpu mannually.