correctly set gpu

pyg-team · Sep 3, 2024 · 35f891a · 35f891a
1 parent c742861
commit 35f891a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/benchmark/pytorch_tabular_benchmark.py b/benchmark/pytorch_tabular_benchmark.py
@@ -69,7 +69,7 @@ def train_tabular_model() -> float:
         auto_lr_find=False,
         batch_size=args.batch_size,
         max_epochs=args.epochs,
-        accelerator=str(device),
+        accelerator='gpu' if device.type == 'cuda' else 'cpu',
     )
 
     optimizer_config = OptimizerConfig()