huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-17 02:58:50 +00:00
Code Issues Packages Projects Releases Wiki Activity

Use shape with holes. (#771)

Browse Source
This commit is contained in:
Laurent Mazare
2023-09-08 08:50:27 +01:00
committed by GitHub
parent 0e250aee4f
commit 989a4807b1
3 changed files with 7 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
candle-examples/examples/segment-anything/model_mask_decoder.rs
View File

@ -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)?;