I think I have solved the issue now. Thanks for the reply. The problem was that I had some variables not initialised as .cuda().
Is there a general principle in general while using the code for gpus. Writing .cuda() everytime is a little annoying