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: