fix device for lpformer example

pyg-team · Jan 17, 2025 · ded9c7c · ded9c7c
1 parent a5e05f2
commit ded9c7c
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/examples/lpformer.py b/examples/lpformer.py
@@ -64,11 +64,11 @@
 model = LPFormer(data.x.size(-1), args.hidden_channels,
                  num_gnn_layers=args.gnn_layers,
                  ppr_thresholds=args.thresholds, gnn_dropout=args.dropout,
-                 transformer_dropout=args.dropout).to(device)
+                 transformer_dropout=args.dropout, gcn_cache=True).to(device)
 
 # Get PPR matrix in sparse format
 ppr_matrix = model.calc_sparse_ppr(data.edge_index, data.num_nodes,
-                                   eps=args.eps)
+                                   eps=args.eps).to(device)
 
 def train_epoch():
     model.train()