Improve device handling and add fallback to CPU

app.py

@@ -6,15 +6,18 @@ from transformer_lens.utils import to_numpy
 import torch
 
 model_name = "gpt2-small"
+# Determine device based on CUDA availability
+device = "cuda" if torch.cuda.is_available() else "cpu"
 model = HookedTransformer.from_pretrained(
     model_name,
-    device="cuda"  
+    device=device  
 )
 
-if not torch.cuda.is_available():
-    raise RuntimeError("This application requires a GPU to run. No GPU was detected.")
-
-print(f"Using GPU: {torch.cuda.get_device_name(0)}")
+# Only print GPU info if using CUDA
+if device == "cuda":
+    print(f"Using GPU: {torch.cuda.get_device_name(0)}")
+else:
+    print("Using CPU")
 
 def get_neuron_acts(text, layer, neuron_index):
     cache = {}