Spaces:

DeepMount00
/

Italian_OCR

Running on Zero

DeepMount00 commited on 29 days ago

Commit

94fc7c9

verified ·

1 Parent(s): 4a7f71c

Update app.py

Files changed (1) hide show

app.py CHANGED Viewed

@@ -29,6 +29,7 @@ model = AutoModelForVision2Seq.from_pretrained(
     MODEL_ID,
     token=hf_token,
     torch_dtype=torch.bfloat16,
 ).to("cuda")  # Ensure model loads on CUDA for Spaces
 @spaces.GPU  # Add spaces.GPU decorator for GPU acceleration

     MODEL_ID,
     token=hf_token,
     torch_dtype=torch.bfloat16,
+    _attn_implementation="flash_attention_2" if DEVICE == "cuda" else "eager",
 ).to("cuda")  # Ensure model loads on CUDA for Spaces
 @spaces.GPU  # Add spaces.GPU decorator for GPU acceleration