huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-18 11:37:11 +00:00
Code Issues Packages Projects Releases Wiki Activity

Do not backprop through argmin/argmax. (#865)

Browse Source
This commit is contained in:
Laurent Mazare
2023-09-15 23:15:40 +02:00
committed by GitHub
parent 3e49f8fce5
commit 635012d770
2 changed files with 9 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
candle-core/src/tensor.rs
View File

@ -666,7 +666,12 @@ impl Tensor {
let storage = self.storage().reduce_op(op, self.layout(), &[dim])?;
let mut dims = self.dims().to_vec();
dims[dim] = 1;
let op = BackpropOp::new1(self, |arg| Op::Reduce(arg, op, dims.to_vec()));
let op = match op {
ReduceOp::Sum | ReduceOp::Min | ReduceOp::Max => {
BackpropOp::new1(self, |arg| Op::Reduce(arg, op, dims.to_vec()))
}
ReduceOp::ArgMin | ReduceOp::ArgMax => BackpropOp::none(),
};
let res = from_storage(storage, dims, op, false);
if keepdim {
Ok(res)