From 33929e05087620c60c00f01fb88ca6ed97f59b81 Mon Sep 17 00:00:00 2001 From: Tauheed Elahee Date: Thu, 12 Mar 2026 21:13:00 -0400 Subject: [PATCH] Use dynamic device count instead of hardcoded num_devices=2 Replace hardcoded num_devices=2 for Mistral with max(1, t.cuda.device_count()) to support single-GPU and CPU-only machines while using all available GPUs when present. --- sae_multid_feature_discovery/generate_feature_occurence_data.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sae_multid_feature_discovery/generate_feature_occurence_data.py b/sae_multid_feature_discovery/generate_feature_occurence_data.py index 63733aa..f5b2a21 100644 --- a/sae_multid_feature_discovery/generate_feature_occurence_data.py +++ b/sae_multid_feature_discovery/generate_feature_occurence_data.py @@ -32,7 +32,7 @@ model_name = "mistral-7b" batch_size = 16 layers_to_evaluate = [8, 16, 24] - num_devices = 2 + num_devices = max(1, t.cuda.device_count()) sae_hidden_size = 65536 else: