Update app.py

app.py

@@ -1,5 +1,7 @@
 import os
-os.environ["XLA_FLAGS"] = "--xla_gpu_force_compilation_parallelism=1"
+# Force CPU-only mode
+os.environ["JAX_PLATFORM_NAME"] = "cpu"
+os.environ["CUDA_VISIBLE_DEVICES"] = ""
 os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"  # Suppress TensorFlow warnings
 
 # Suppress JAX compilation warnings
@@ -176,13 +178,8 @@ class SimpleMAXIMPredictor:
         
         try:
             import jax
-            # Check if GPU is available
-            gpu_available = len(jax.devices('gpu')) > 0
-            device_info = f"Using device: {jax.default_backend()}"
-            if gpu_available:
-                device_info += f" (GPU: {jax.devices('gpu')[0]})"
-            else:
-                device_info += " (CPU fallback)"
+            # Force CPU mode - device info
+            device_info = f"Using device: {jax.default_backend()} (CPU-only mode)"
             print(device_info)
             
             # Get model and parameters