Auto set device

2024-04-17 16:42:55 +02:00 · 2024-04-17 16:42:55 +02:00 · a870907266
parent 75d34548a5
commit a870907266
1 changed files with 2 additions and 2 deletions
--- a/tests/utils.py
+++ b/tests/utils.py
@ -1,6 +1,6 @@
-import os
+import torch

 # Pass this as the first argument to init_hydra_config.
 DEFAULT_CONFIG_PATH = "lerobot/configs/default.yaml"

-DEVICE = os.environ.get('LEROBOT_TESTS_DEVICE', "cuda")
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"