diff options
Diffstat (limited to 'candle-nn/src')
| -rw-r--r-- | candle-nn/src/ops.rs | 3 |
1 files changed, 2 insertions, 1 deletions
diff --git a/candle-nn/src/ops.rs b/candle-nn/src/ops.rs index aaec8b56..fdd67142 100644 --- a/candle-nn/src/ops.rs +++ b/candle-nn/src/ops.rs @@ -238,7 +238,8 @@ impl candle::CustomOp1 for SoftmaxLastDim { &output, ) .unwrap(); - let newstorage = candle::MetalStorage::new(output, device.clone(), storage.dtype()); + let newstorage = + candle::MetalStorage::new(output, device.clone(), elem_count, storage.dtype()); Ok((newstorage, layout.shape().clone())) } } |