import torch

rapidsai · Aug 22, 2024 · a22ce90 · a22ce90
1 parent 1fcd982
commit a22ce90
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/python/cugraph/cugraph/gnn/data_loading/dist_sampler.py b/python/cugraph/cugraph/gnn/data_loading/dist_sampler.py
@@ -236,6 +236,8 @@ def __sample_from_nodes_func(
         random_state: int,
         assume_equal_input_size: bool,
     ) -> Union[None, Iterator[Tuple[Dict[str, "torch.Tensor"], int, int]]]:
+        torch = import_optional("torch")
+
         current_batches = torch.arange(
             batch_id_start + call_id * batches_per_call,
             batch_id_start