Change to bfloat16

app.py

@@ -301,7 +301,7 @@ def run_model(target_dir, model) -> dict:
 
     # 3. Infer
     print("Running model inference...")
-    dtype = torch.float16
+    dtype = torch.bfloat16
     with torch.no_grad():
         with torch.amp.autocast('cuda', dtype=dtype):
             predictions = model(imgs[None]) # Add batch dimension