Update app.py

app.py

@@ -4,13 +4,14 @@ import os
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
-# Set environment variables for GPU usage and memory allocation
-os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'
-torch.cuda.empty_cache()
-torch.cuda.set_per_process_memory_fraction(0.5)  # Adjust the fraction as needed
-
-# Define device
-device = "cuda"  # The device to load the model onto
+# Check if CUDA is available and set device accordingly
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Set environment variables for GPU usage and memory allocation if CUDA is available
+if device == "cuda":
+    os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'
+    torch.cuda.empty_cache()
+    torch.cuda.set_per_process_memory_fraction(0.8)  # Adjust the fraction as needed
 
 # System message (placeholder, adjust as needed)
 system_message = ""