Update app.py

app.py

@@ -167,9 +167,9 @@ processor_s = AutoProcessor.from_pretrained(MODEL_ID_S, trust_remote_code=True)
 model_s = AutoModelForCausalLM.from_pretrained(
     MODEL_ID_S,
     _attn_implementation='flash_attention_2',
+    torch_dtype=torch.bfloat16,
     trust_remote_code=True,
-    use_safetensors=True
-).eval().to(device).to(torch.bfloat16)
+).eval()
 
 
 @spaces.GPU