Fix

Signed-off-by: Serge Panev <[email protected]>
pyg-team · Jul 7, 2023 · 84c7a6d · 84c7a6d
1 parent da89da4
commit 84c7a6d
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torch_geometric/sampler/neighbor_sampler.py b/torch_geometric/sampler/neighbor_sampler.py
@@ -162,7 +162,7 @@ def hook_func_with_nvtx(func_name):
                 def hooked_func(*args, **kwargs):
                     nvtx_handle = torch.cuda.nvtx.range_start(
                         f"[Sampler] {func_name} for {self}")
-                    ret = func(**kwargs)
+                    ret = func(*args, **kwargs)
                     torch.cuda.nvtx.range_end(nvtx_handle)
                     return ret