Spaces:
Running
on
Zero
Running
on
Zero
Reduce VRAM consumption by swapping `cuda()` and `to(torch.bfloat16)`
#2
by
mingyi456
- opened
app.py
CHANGED
|
@@ -43,7 +43,7 @@ def process_ocr_task(image, model_size, task_type, ref_text):
|
|
| 43 |
return "Please upload an image first.", None
|
| 44 |
|
| 45 |
print("🚀 Moving model to GPU...")
|
| 46 |
-
model_gpu = model.
|
| 47 |
print("✅ Model is on GPU.")
|
| 48 |
|
| 49 |
with tempfile.TemporaryDirectory() as output_path:
|
|
|
|
| 43 |
return "Please upload an image first.", None
|
| 44 |
|
| 45 |
print("🚀 Moving model to GPU...")
|
| 46 |
+
model_gpu = model.to(torch.bfloat16).cuda()
|
| 47 |
print("✅ Model is on GPU.")
|
| 48 |
|
| 49 |
with tempfile.TemporaryDirectory() as output_path:
|