Fix torch bloat16 -> numpy float32 conversion for compile max-autotune (

#96)
pytorch-labs · Dec 7, 2023 · ceaae4b · ceaae4b
1 parent 97a69cf
commit ceaae4b
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/segment_anything_fast/predictor.py b/segment_anything_fast/predictor.py
@@ -162,8 +162,8 @@ def predict(
         )
 
         masks_np = masks[0].detach().cpu().numpy()
-        iou_predictions_np = iou_predictions[0].detach().cpu().numpy()
-        low_res_masks_np = low_res_masks[0].detach().cpu().numpy()
+        iou_predictions_np = iou_predictions[0].detach().cpu().float().numpy()
+        low_res_masks_np = low_res_masks[0].detach().cpu().float().numpy()
         return masks_np, iou_predictions_np, low_res_masks_np
 
     @torch.no_grad()