From f26cfc96beaaf05c4bbf1b5d35f99047772b9bed Mon Sep 17 00:00:00 2001
From: jianyuan <jianyuan@meta.com>
Date: Tue, 18 Mar 2025 04:12:19 +0000
Subject: [PATCH] solve the problem that pytorch will ignore dtype =
 torch.bfloat16 if bfloat not supported

---
 README.md | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 6dc106a..087deef 100644
--- a/README.md
+++ b/README.md
@@ -47,7 +47,8 @@ from vggt.models.vggt import VGGT
 from vggt.utils.load_fn import load_and_preprocess_images
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
-dtype = torch.bfloat16  # or torch.float16
+# bfloat16 is supported on Ampere GPUs (Compute Capability 8.0+) 
+dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] >= 8 else torch.float16
 
 # Initialize the model and load the pretrained weights.
 # This will automatically download the model weights the first time it's run, which may take a while.