From 313192670d0b6867c9620456db4c130503ed3ff9 Mon Sep 17 00:00:00 2001 From: Alexandria Barghi Date: Wed, 28 Aug 2024 08:54:44 -0700 Subject: [PATCH] revert change to dgl sampler --- python/cugraph-dgl/cugraph_dgl/dataloading/neighbor_sampler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/cugraph-dgl/cugraph_dgl/dataloading/neighbor_sampler.py b/python/cugraph-dgl/cugraph_dgl/dataloading/neighbor_sampler.py index 1b5b1eb5159..a2ac18a9585 100644 --- a/python/cugraph-dgl/cugraph_dgl/dataloading/neighbor_sampler.py +++ b/python/cugraph-dgl/cugraph_dgl/dataloading/neighbor_sampler.py @@ -194,7 +194,7 @@ def sample( if g.is_homogeneous: indices = torch.concat(list(indices)) - reader = ds.sample_from_nodes(indices, batch_size=batch_size) + reader = ds.sample_from_nodes(indices.long(), batch_size=batch_size) return HomogeneousSampleReader(reader, self.output_format, self.edge_dir) raise ValueError(