Fix the problem that some old gpus do not support bfloat16
This commit is contained in:
+3
-1
@@ -68,8 +68,10 @@ def run_model(target_dir, model) -> dict:
|
||||
|
||||
# Run 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)
|
||||
|
||||
# Convert pose encoding to extrinsic and intrinsic matrices
|
||||
|
||||
Reference in New Issue
Block a user