Update app.py

app.py

@@ -21,8 +21,6 @@ def preload_models(model_choices):
                 torch_dtype=torch.bfloat16,
                 trust_remote_code=True,
                 token=os.environ.get("token"),
-                device_map="cpu",
-                low_cpu_mem_usage=True
             )
             
             # Load tokenizer
@@ -46,7 +44,7 @@ def get_model_pipeline(model_name):
         raise ValueError(f"Model {model_name} not found in preloaded models")
     
     # Move model to GPU
-    model = LOADED_MODELS[model_name].to("cuda")
+    model = LOADED_MODELS[model_name]
     tokenizer = LOADED_TOKENIZERS[model_name]
     
     # Create pipeline with the GPU model
@@ -55,6 +53,7 @@ def get_model_pipeline(model_name):
         model=model,
         tokenizer=tokenizer,
         torch_dtype=torch.bfloat16,
+        device="cuda"
     )
     
     return pipe, model