Update app.py

app.py

@@ -28,6 +28,7 @@ def replace_single_quotes(text):
 
 class PromptEnhancerV2:
     def __init__(self, models_root_path, device_map="cuda", torch_dtype="bfloat16"):#auto
+        device_map = "cuda:0"
         if not logging.getLogger(__name__).handlers:
             logging.basicConfig(level=logging.INFO)
         self.logger = logging.getLogger(__name__)
@@ -57,7 +58,7 @@ class PromptEnhancerV2:
         temperature=0.0,
         top_p=1.0,
         max_new_tokens=2048,
-        device="cuda",
+        device="cuda:0",
     ):
         org_prompt_cot = prompt_cot
         try:
@@ -122,7 +123,7 @@ class PromptEnhancerV2:
         temperature=0.1,
         top_p=1.0,
         max_new_tokens=2048,
-        device="cuda",
+        device="cuda:0",
     ):
         org_prompt_cot = prompt_cot