diff --git a/examples/hetero/bipartite_sage_unsup.py b/examples/hetero/bipartite_sage_unsup.py index 319bf3a6edb59..fa99d54676d98 100644 --- a/examples/hetero/bipartite_sage_unsup.py +++ b/examples/hetero/bipartite_sage_unsup.py @@ -13,8 +13,14 @@ from torch_geometric.loader import LinkNeighborLoader from torch_geometric.nn import SAGEConv from torch_geometric.utils.convert import to_scipy_sparse_matrix - -device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') +from torch_geometric.testing.device import is_xpu_avaliable + +if torch.cuda.is_available(): + device = torch.device('cuda:0') +elif is_xpu_avaliable(): + device = torch.device('xpu:0') +else: + device = torch.device('cpu') path = osp.join(osp.dirname(osp.realpath(__file__)), '../../data/Taobao') dataset = Taobao(path) diff --git a/examples/hetero/dmgi_unsup.py b/examples/hetero/dmgi_unsup.py index 51dd7c1421c01..8238685e053b8 100644 --- a/examples/hetero/dmgi_unsup.py +++ b/examples/hetero/dmgi_unsup.py @@ -12,6 +12,7 @@ import torch_geometric.transforms as T from torch_geometric.datasets import IMDB from torch_geometric.nn import GCNConv +from torch_geometric.testing.device import is_xpu_avaliable path = osp.join(osp.dirname(osp.realpath(__file__)), '../../data/IMDB') dataset = IMDB(path) @@ -74,7 +75,13 @@ def loss(self, pos_hs, neg_hs, summaries): model = DMGI(data['movie'].num_nodes, data['movie'].x.size(-1), out_channels=64, num_relations=len(data.edge_types)) -device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + +if torch.cuda.is_available(): + device = torch.device('cuda') +elif is_xpu_avaliable(): + device = torch.device('xpu') +else: + device = torch.device('cpu') data, model = data.to(device), model.to(device) optimizer = Adam(model.parameters(), lr=0.0005, weight_decay=0.0001) diff --git a/examples/hetero/han_imdb.py b/examples/hetero/han_imdb.py index 22ee0ced291bd..bdf2498650a18 100644 --- a/examples/hetero/han_imdb.py +++ b/examples/hetero/han_imdb.py @@ -8,6 +8,7 @@ import torch_geometric.transforms as T from torch_geometric.datasets import IMDB from torch_geometric.nn import HANConv +from torch_geometric.testing.device import is_xpu_avaliable path = osp.join(osp.dirname(osp.realpath(__file__)), '../../data/IMDB') metapaths = [[('movie', 'actor'), ('actor', 'movie')], @@ -34,7 +35,12 @@ def forward(self, x_dict, edge_index_dict): model = HAN(in_channels=-1, out_channels=3) -device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') +if torch.cuda.is_available(): + device = torch.device('cuda') +elif is_xpu_avaliable(): + device = torch.device('xpu') +else: + device = torch.device('cpu') data, model = data.to(device), model.to(device) with torch.no_grad(): # Initialize lazy modules. diff --git a/examples/hetero/hetero_conv_dblp.py b/examples/hetero/hetero_conv_dblp.py index 51820b10a5328..70dfdc2a9e9ea 100644 --- a/examples/hetero/hetero_conv_dblp.py +++ b/examples/hetero/hetero_conv_dblp.py @@ -6,6 +6,7 @@ import torch_geometric.transforms as T from torch_geometric.datasets import DBLP from torch_geometric.nn import HeteroConv, Linear, SAGEConv +from torch_geometric.testing.device import is_xpu_avaliable path = osp.join(osp.dirname(osp.realpath(__file__)), '../../data/DBLP') # We initialize conference node features with a single one-vector as feature: @@ -37,7 +38,12 @@ def forward(self, x_dict, edge_index_dict): model = HeteroGNN(data.metadata(), hidden_channels=64, out_channels=4, num_layers=2) -device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') +if torch.cuda.is_available(): + device = torch.device('cuda') +elif is_xpu_avaliable(): + device = torch.device('xpu') +else: + device = torch.device('cpu') data, model = data.to(device), model.to(device) with torch.no_grad(): # Initialize lazy modules. diff --git a/examples/hetero/hetero_link_pred.py b/examples/hetero/hetero_link_pred.py index 4acc451e36e6e..9f7699395bb9c 100644 --- a/examples/hetero/hetero_link_pred.py +++ b/examples/hetero/hetero_link_pred.py @@ -8,6 +8,7 @@ import torch_geometric.transforms as T from torch_geometric.datasets import MovieLens from torch_geometric.nn import SAGEConv, to_hetero +from torch_geometric.testing.device import is_xpu_avaliable parser = argparse.ArgumentParser() parser.add_argument('--use_weighted_loss', action='store_true', @@ -18,6 +19,8 @@ device = torch.device('cuda') elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available(): device = torch.device('mps') +elif is_xpu_avaliable(): + device = torch.device('xpu') else: device = torch.device('cpu') diff --git a/examples/hetero/hgt_dblp.py b/examples/hetero/hgt_dblp.py index f21149bf3f125..645eadcea54c1 100644 --- a/examples/hetero/hgt_dblp.py +++ b/examples/hetero/hgt_dblp.py @@ -6,6 +6,7 @@ import torch_geometric.transforms as T from torch_geometric.datasets import DBLP from torch_geometric.nn import HGTConv, Linear +from torch_geometric.testing.device import is_xpu_avaliable path = osp.join(osp.dirname(osp.realpath(__file__)), '../../data/DBLP') # We initialize conference node features with a single one-vector as feature: @@ -43,7 +44,12 @@ def forward(self, x_dict, edge_index_dict): model = HGT(hidden_channels=64, out_channels=4, num_heads=2, num_layers=1) -device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') +if torch.cuda.is_available(): + device = torch.device('cuda') +elif is_xpu_avaliable(): + device = torch.device('xpu') +else: + device = torch.device('cpu') data, model = data.to(device), model.to(device) with torch.no_grad(): # Initialize lazy modules. diff --git a/examples/hetero/hierarchical_sage.py b/examples/hetero/hierarchical_sage.py index 8b32e5ae67a41..96b939f0b9f8f 100644 --- a/examples/hetero/hierarchical_sage.py +++ b/examples/hetero/hierarchical_sage.py @@ -15,7 +15,12 @@ parser.add_argument('--use-sparse-tensor', action='store_true') args = parser.parse_args() -device = args.device if torch.cuda.is_available() else 'cpu' +if args.device=='cuda' and torch.cuda.is_available(): + device = 'cuda' +elif args.device=='xpu' and is_xpu_avaliable(): + device = 'xpu' +else: + device = 'cpu' transforms = [T.ToUndirected(merge=True)] if args.use_sparse_tensor: diff --git a/examples/hetero/metapath2vec.py b/examples/hetero/metapath2vec.py index 3d8801cea2112..25c4fbcdf21ab 100644 --- a/examples/hetero/metapath2vec.py +++ b/examples/hetero/metapath2vec.py @@ -6,6 +6,7 @@ from torch_geometric.datasets import AMiner from torch_geometric.nn import MetaPath2Vec +from torch_geometric.testing.device import is_xpu_avaliable path = osp.join(osp.dirname(osp.realpath(__file__)), '../../data/AMiner') dataset = AMiner(path) @@ -18,7 +19,12 @@ ('paper', 'written_by', 'author'), ] -device = 'cuda' if torch.cuda.is_available() else 'cpu' +if torch.cuda.is_available(): + device = 'cuda' +elif is_xpu_avaliable(): + device = 'xpu' +else: + device = 'cpu' model = MetaPath2Vec(data.edge_index_dict, embedding_dim=128, metapath=metapath, walk_length=50, context_size=7, walks_per_node=5, num_negative_samples=5,