Solved by eqy in post #2 
              
              
                Have you checked if x,y (the test data here) are also on GPU or if they are on CPU? 
The training code has 
X,y=X.to(device),y.to(device);

so something similar might be needed here.