Fix breaking torch.device(0) incompatibility for torch 2.5 (#151)

Author: nshah171

benchmarks/api/bench_dist_neighbor_loader.py

@@ -123,7 +123,7 @@
     drop_last=True,
     with_edge=with_edge,
     collect_features=collect_features,
-    to_device=torch.device(0),
+    to_device=torch.device('cuda:0'),
     worker_options=glt.distributed.MpDistSamplingWorkerOptions(
       num_workers=sampling_nprocs,
       worker_devices=[torch.device('cuda', i % device_count) for i in range(sampling_nprocs)],

benchmarks/api/bench_feature.py

@@ -36,7 +36,7 @@ def test_glt_ogbnproducts(split_ratio):
   csr_topo = glt.data.Topology(dataset[0].edge_index)
 
   g = glt.data.Graph(csr_topo, 'CUDA', device=0)
-  device = torch.device(0)
+  device = torch.device('cuda:0')
   sampler = glt.sampler.NeighborSampler(g, [15, 10, 5], device=device)
 
   cpu_tensor, id2index = glt.data.sort_by_in_degree(

benchmarks/api/bench_sampler.py

@@ -39,7 +39,7 @@ def test_glt_ogbnproducts(mode='GPU'):
                                              shuffle=True)
   csr_topo = glt.data.Topology(dataset[0].edge_index)
   g = glt.data.Graph(csr_topo, graph_mode, device=0)
-  device = torch.device(0)
+  device = torch.device('cuda:0')
   sampler = glt.sampler.NeighborSampler(g, [15, 10, 5], device=device)
   total_time = 0
   sampled_edges = 0

docs/get_started/node_class.md

@@ -85,7 +85,7 @@ train_loader = glt.loader.NeighborLoader(glt_dataset,
                                          batch_size=1024,
                                          shuffle=True,
                                          drop_last=True,
-                                         device=torch.device(0),
+                                         device=torch.device('cuda:0'),
                                          as_pyg_v1=True)
 ```
 > **Note**

docs/tutorial/dist.md

@@ -116,7 +116,7 @@ random_partitioner = glt.partition.RandomPartitioner(
   edge_feat=None,
   edge_assign_strategy='by_src', # assign graph edges by src node.
   chunk_size=10000, # chunk size for node assignment
-  device=torch.device(0) # device used for partitioning
+  device=torch.device('cuda:0') # device used for partitioning
 )
 random_partitioner.partition()
 ```
@@ -170,7 +170,7 @@ freq_partitioner = glt.partition.FrequencyPartitioner(
   edge_assign_strategy='by_src', # assign graph edges by src node.
   chunk_size=10000, # chunk size for node assignment
   cache_ratio=0.2, # cache 20% hottest graph nodes
-  device=torch.device(0) # device used for partitioning
+  device=torch.device('cuda:0') # device used for partitioning
 )
 freq_partitioner.partition()
 ```
@@ -507,7 +507,7 @@ train_loader = glt.distributed.DistNeighborLoader(
   # Set to true to collect node features for sampled subgraphs.
   collect_features=True,
   # All sampled results will be moved to this device.
-  to_device=torch.device(0),
+  to_device=torch.device('cuda:0'),
   # Use `MpDistSamplingWorkerOptions`.
   worker_options=mp_options
 )
@@ -589,7 +589,7 @@ train_loader = glt.distributed.DistNeighborLoader(
   input_nodes=train_idx,
   batch_size=1024,
   collect_features=True,
-  to_device=torch.device(0),
+  to_device=torch.device('cuda:0'),
   worker_options=collocated_options
 )
 ```
@@ -645,7 +645,7 @@ train_loader = glt.distributed.DistNeighborLoader(
   input_nodes=train_idx,
   batch_size=1024,
   collect_features=True,
-  to_device=torch.device(0),
+  to_device=torch.device('cuda:0'),
   worker_options=remote_options
 )
 ```

docs/tutorial/graph_ops.md

@@ -38,7 +38,7 @@ class NeighborSampler(BaseSampler):
   def __init__(self,
                graph: Union[Graph, Dict[str, Graph]],
                num_neighbors: NumNeighbors,
-               device: torch.device=torch.device('cuda', 0),
+               device: torch.device=torch.device('cuda:0'),
                with_edge: bool=False,
                strategy: str = 'random'):
 ```

examples/train_sage_ogbn_products.py

@@ -106,7 +106,7 @@ def inference(self, x_all):
 )
 glt_dataset.init_node_labels(node_label_data=data.y)
 
-device = torch.device(0)
+device = torch.device('cuda:0')
 # graphlearn_torch NeighborLoader
 train_loader = glt.loader.NeighborLoader(glt_dataset,
                                          [15, 10, 5],

graphlearn_torch/python/loader/link_loader.py

@@ -104,7 +104,7 @@ def __init__(
     edge_label_index: InputEdges = None,
     edge_label: Optional[torch.Tensor] = None,
     neg_sampling: Optional[NegativeSampling] = None,
-    device: torch.device = torch.device(0),
+    device: torch.device = torch.device('cuda:0'),
     edge_dir: Literal['out', 'in'] = 'out',
     **kwargs,
   ):

graphlearn_torch/python/loader/link_neighbor_loader.py

@@ -122,7 +122,7 @@ def __init__(
     shuffle: bool = False,
     drop_last: bool = False,
     strategy: str = "random",
-    device: torch.device = torch.device(0),
+    device: torch.device = torch.device('cuda:0'),
     seed: Optional[int] = None,
     **kwargs,
   ):

graphlearn_torch/python/loader/neighbor_loader.py

@@ -68,7 +68,7 @@ def __init__(
     with_edge: bool = False,
     with_weight: bool = False,
     strategy: str = 'random',
-    device: torch.device = torch.device(0),
+    device: torch.device = torch.device('cuda:0'),
     as_pyg_v1: bool = False,
     seed: Optional[int] = None,
     **kwargs

graphlearn_torch/python/loader/node_loader.py

@@ -56,7 +56,7 @@ def __init__(
     data: Dataset,
     node_sampler: BaseSampler,
     input_nodes: InputNodes,
-    device: torch.device = torch.device(0),
+    device: torch.device = torch.device('cuda:0'),
     **kwargs
   ):
     self.data = data

graphlearn_torch/python/loader/subgraph_loader.py

@@ -62,7 +62,7 @@ def __init__(
     drop_last: bool = False,
     with_edge: bool = False,
     strategy: str = 'random',
-    device: torch.device = torch.device(0),
+    device: torch.device = torch.device('cuda:0'),
     seed: Optional[int] = None,
     **kwargs
   ):

graphlearn_torch/python/sampler/neighbor_sampler.py

@@ -41,7 +41,7 @@ class NeighborSampler(BaseSampler):
   def __init__(self,
                graph: Union[Graph, Dict[EdgeType, Graph]],
                num_neighbors: Optional[NumNeighbors] = None,
-               device: torch.device=torch.device('cuda', 0),
+               device: torch.device=torch.device('cuda:0'),
                with_edge: bool=False,
                with_neg: bool=False,
                with_weight: bool=False,