feat: use bfloat16

app.py

@@ -16,6 +16,7 @@ from src.plot_utils import export_mask
 @spaces.GPU()
 def predict(model_choice, annotations: Dict[str, Any]):
     # device = "cuda" if torch.cuda.is_available() else "cpu"
+    torch.autocast(device_type="cuda", dtype=torch.bfloat16).__enter__()
     sam2_model = load_model(
         variant=model_choice,
         ckpt_path=f"assets/checkpoints/sam2_hiera_{model_choice}.pt",