fix eval device

Lightning-AI · Apr 23, 2024 · 973a65a · 973a65a
1 parent bd8cec9
commit 973a65a
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/litgpt/eval/evaluate.py b/litgpt/eval/evaluate.py
@@ -72,6 +72,9 @@ def convert_and_evaluate(
         )
         return
 
+    if device is None:
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+
     checkpoint_dir = Path(checkpoint_dir)
 
     if out_dir is None:

diff --git a/tests/test_evaluate.py b/tests/test_evaluate.py
@@ -39,7 +39,7 @@ def test_evaluate_script(tmp_path, monkeypatch):
     fn_kwargs = dict(
         checkpoint_dir=tmp_path,
         out_dir=tmp_path / "out_dir",
-        device="cpu",
+        device=None,
         dtype=torch.float32,
         limit=5,
         tasks="mathqa"