fix: fix pyg remote backend ut (#147)

alibaba · Nov 1, 2024 · 36ce42b · 36ce42b
1 parent 2034837
commit 36ce42b
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/graphlearn_torch/python/distributed/dist_server.py b/graphlearn_torch/python/distributed/dist_server.py
@@ -20,7 +20,7 @@
 import warnings
 
 import torch
-from graphscope.learning.graphlearn_torch.partition.base import PartitionBook
+from ..partition import PartitionBook
 
 from ..channel import ShmChannel, QueueTimeoutError
 from ..sampler import NodeSamplerInput, EdgeSamplerInput, SamplingConfig, RemoteSamplerInput
@@ -95,7 +95,7 @@ def get_node_partition_id(self, node_type, index):
 
   def get_node_feature(self, node_type, index):
     feature = self.dataset.get_node_feature(node_type)
-    return feature[index]
+    return feature[index].cpu()
 
   def get_tensor_size(self, node_type):
     feature = self.dataset.get_node_feature(node_type)