From aeccf2ef2f5ff871d4ed4fda65b1ca3c77a6ea2a Mon Sep 17 00:00:00 2001
From: Jianyuan Wang <37647961+jytime@users.noreply.github.com>
Date: Wed, 19 Mar 2025 03:53:36 +0000
Subject: [PATCH] Fix the problem that some old gpus do not support bfloat16

---
 demo_gradio.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/demo_gradio.py b/demo_gradio.py
index da70f0a..8dcef99 100644
--- a/demo_gradio.py
+++ b/demo_gradio.py
@@ -68,8 +68,10 @@ def run_model(target_dir, model) -> dict:
 
     # Run inference
     print("Running inference...")
+    dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] >= 8 else torch.float16
+
     with torch.no_grad():
-        with torch.cuda.amp.autocast(dtype=torch.bfloat16):
+        with torch.cuda.amp.autocast(dtype=dtype):
             predictions = model(images)
 
     # Convert pose encoding to extrinsic and intrinsic matrices