From 7c71ff01118d579a99083d0535989fb3e02e1abf Mon Sep 17 00:00:00 2001 From: Alexandria Barghi Date: Wed, 13 Mar 2024 13:17:50 -0700 Subject: [PATCH] remove cuda check --- python/cugraph-pyg/cugraph_pyg/data/cugraph_store.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/python/cugraph-pyg/cugraph_pyg/data/cugraph_store.py b/python/cugraph-pyg/cugraph_pyg/data/cugraph_store.py index df16fc9fd6c..354eea8ee6b 100644 --- a/python/cugraph-pyg/cugraph_pyg/data/cugraph_store.py +++ b/python/cugraph-pyg/cugraph_pyg/data/cugraph_store.py @@ -336,9 +336,7 @@ def __del__(self): def __make_offsets(self, input_dict): offsets = {} offsets["stop"] = [input_dict[v] for v in sorted(input_dict.keys())] - offsets["stop"] = torch.tensor(offsets["stop"]) - if torch.has_cuda: - offsets["stop"] = offsets["stop"].cuda() + offsets["stop"] = torch.tensor(offsets["stop"]).cuda() cumsum = offsets["stop"].cumsum(0) offsets["start"] = cumsum - offsets["stop"]