Solved by albanD in post #2 
              
              
                You need to do QPs = QPs.to(output.device) for it to work as to does not change the Tensor inplace. 
Also you can try QPs = output.new(QPs, device=output.device).