Update app.py

app.py

@@ -228,8 +228,7 @@ def build_models(model_type, config, enable_optimization=False):
     stdit = AutoModel.from_pretrained(
         HF_STDIT_MAP[model_type],
         enable_flash_attn=enable_optimization,
-        trust_remote_code=True,
-        device_map="auto",
+        trust_remote_code=True
     ).cuda()
 
     # build scheduler