Use shape with holes. (#771)

1 files changed, 1 insertions, 1 deletions

diff --git a/candle-examples/examples/segment-anything/model_mask_decoder.rs b/candle-examples/examples/segment-anything/model_mask_decoder.rs
index 1ef46eeb..acbfeeea 100644
--- a/candle-examples/examples/segment-anything/model_mask_decoder.rs
+++ b/candle-examples/examples/segment-anything/model_mask_decoder.rs
@@ -214,7 +214,7 @@ impl MaskDecoder {
         let hyper_in = Tensor::stack(hyper_in_list.as_slice(), 1)?;
         let (b, c, h, w) = upscaled_embedding.dims4()?;
         let masks = hyper_in.matmul(&upscaled_embedding.reshape((b, c, h * w))?)?;
-        let masks = masks.reshape((b, masks.elem_count() / b / h / w, h, w))?;
+        let masks = masks.reshape((b, (), h, w))?;
 
         // Generate mask quality predictions.
         let iou_pred = self.iou_prediction_head.forward(&iou_token_out)?;