Fix the problem that some old gpus do not support bfloat16

This commit is contained in:
Jianyuan Wang
2025-03-19 03:54:29 +00:00
committed by GitHub
parent aeccf2ef2f
commit fa5911cb6b
+4 -2
View File
@@ -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...")