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

mawanda-jun · web-flow · commit ceaae4baabd0 · 2023-12-07T13:00:52.000-05:00
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()

Original file line number	Diff line number	Diff line change
`@@ -162,8 +162,8 @@ def predict(`
`162`	`162`	`)`
`163`	`163`
`164`	`164`	`masks_np = masks[0].detach().cpu().numpy()`
`165`		`- iou_predictions_np = iou_predictions[0].detach().cpu().numpy()`
`166`		`- low_res_masks_np = low_res_masks[0].detach().cpu().numpy()`
	`165`	`+ iou_predictions_np = iou_predictions[0].detach().cpu().float().numpy()`
	`166`	`+ low_res_masks_np = low_res_masks[0].detach().cpu().float().numpy()`
`167`	`167`	`return masks_np, iou_predictions_np, low_res_masks_np`
`168`	`168`
`169`	`169`	`@torch.no_grad()`