Update app.py

app.py

@@ -59,13 +59,13 @@ cls_model_name = "saiteki-kai/QA-DeBERTa-v3-large-binary-3"
 model = transformers.pipeline(
     model=chat_model_name,
     model_kwargs={"dtype": torch.bfloat16} | ({"attn_implementation": "flash_attention_2"} if is_flash_attn_2_available() else {}),
-    device_map="cuda",
+    device=device,
 )
 
 classifier = transformers.pipeline(
     model=cls_model_name, 
     model_kwargs={"dtype": torch.bfloat16}, 
-    device_map="cuda"
+    device=device
 )
 
 unsafe_idx = classifier.model.config.label2id["unsafe"]