Improving the MLP example. (#6593)

examples/pytorch/ogb/ogbn-products/mlp/mlp.py

@@ -99,6 +99,7 @@ def train(
 
     preds = torch.zeros(labels.shape[0], n_classes)
 
+    with dataloader.enable_cpu_affinity():
         for _input_nodes, output_nodes, subgraphs in dataloader:
             subgraphs = [b.to(device) for b in subgraphs]
             new_train_idx = list(range(len(output_nodes)))
@@ -143,6 +144,7 @@ def evaluate(
     eval_times = 1  # Due to the limitation of memory capacity, we calculate the average of logits 'eval_times' times.
 
     for _ in range(eval_times):
+        with dataloader.enable_cpu_affinity():
             for _input_nodes, output_nodes, subgraphs in dataloader:
                 subgraphs = [b.to(device) for b in subgraphs]