solve the problem that pytorch will ignore dtype = torch.bfloat16 if bfloat not supported

2025-03-18 04:12:19 +00:00
parent 56e693172d
commit f26cfc96be
1 changed files with 2 additions and 1 deletions
@@ -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.