Fix the problem that some old gpus do not support bfloat16
This commit is contained in:
+4
-2
@@ -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...")
|
||||
|
||||
Reference in New Issue
Block a user