[Example][Bugfix] graphsage node classification example (#4260)

* Fix node_classification.py

* Minor update
Co-authored-by: nv-dlasalle <63612878+nv-dlasalle@users.noreply.github.com>

examples/pytorch/graphsage/node_classification.py

@@ -72,9 +72,9 @@ def evaluate(model, graph, dataloader):
 def layerwise_infer(device, graph, nid, model, batch_size):
     model.eval()
     with torch.no_grad():
-        pred = model.inference(graph, device, batch_size).to(device)
+        pred = model.inference(graph, device, batch_size) # pred in buffer_device
         pred = pred[nid]
-        label = graph.ndata['label'][nid]
+        label = graph.ndata['label'][nid].to(pred.device)
         return MF.accuracy(pred, label)
 
 def train(args, device, g, dataset, model):
@@ -141,5 +141,5 @@ if __name__ == '__main__':
 
     # test the model
     print('Testing...')
-    acc = layerwise_infer(device, g, dataset.test_idx.to(device), model, batch_size=4096)
+    acc = layerwise_infer(device, g, dataset.test_idx, model, batch_size=4096)
     print("Test Accuracy {:.4f}".format(acc.item()))