Update neighbor_sampler.py

rllm-team · Sep 9, 2024 · 118a0da · 118a0da
1 parent 741d7cc
commit 118a0da
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/rllm/sampler/neighbor_sampler.py b/rllm/sampler/neighbor_sampler.py
@@ -3,7 +3,9 @@
 
 import torch
 from torch import Tensor
+from torch_sparse import SparseTensor
 
+from rllm.utils.sparse import is_torch_sparse_tensor
 
 class NeighborSampler(torch.utils.data.DataLoader):
     r"""The neighbor sampler from the `"Inductive Representation Learning on