From 4f3967dba8c69edffcaecf76d905652074a67a3b Mon Sep 17 00:00:00 2001
From: Adam Karvonen <adam.karvonen@gmail.com>
Date: Sun, 24 Nov 2024 10:08:04 +0000
Subject: [PATCH] Use expandable segments to reduce memory usage

---
 evals/core/main.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/evals/core/main.py b/evals/core/main.py
index 515306e..dd0b7ad 100644
--- a/evals/core/main.py
+++ b/evals/core/main.py
@@ -902,7 +902,7 @@ def multiple_evals(
     verbose: bool = False,
     dtype: str = "float32",
 ) -> List[Dict[str, Any]]:
-    device = "cuda" if torch.cuda.is_available() else "cpu"
+    device = general_utils.setup_environment()
     assert len(filtered_saes) > 0, "No SAEs to evaluate"
 
     eval_results = []