Upload app.py

app.py

@@ -21,9 +21,14 @@ if not torch.cuda.is_available():
 
 if torch.cuda.is_available():
     model_id = "deepseek-ai/deepseek-r1"
-    model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, device_map="auto")
-    tokenizer = AutoTokenizer.from_pretrained(model_id)
-    tokenizer.use_default_system_prompt = False
+if torch.cuda.is_available():
+    model = AutoModelForCausalLM.from_pretrained(
+        model_id, torch_dtype=torch.bfloat16, device_map="auto"
+    )
+else:
+    model = AutoModelForCausalLM.from_pretrained(model_id)
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+tokenizer.use_default_system_prompt = False
 
 @spaces.GPU
 def generate(message: str,