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

---
 demo_viser.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/demo_viser.py b/demo_viser.py
index eef6fc8..c52726b 100644
--- a/demo_viser.py
+++ b/demo_viser.py
@@ -373,9 +373,11 @@ def main():
     images = load_and_preprocess_images(image_names).to(device)
     print(f"Preprocessed images shape: {images.shape}")
 
-    print("Running 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)
 
     print("Converting pose encoding to extrinsic and intrinsic matrices...")