Add binary and ternary custom ops. (#217)

This commit is contained in:
Laurent Mazare
2023-07-21 18:29:50 +02:00
committed by GitHub
parent 4a100875bf
commit 5cc843550d
5 changed files with 209 additions and 10 deletions

View File

@ -94,10 +94,10 @@ impl CustomOp1 for EluWithBackward {
self.0.cpu_fwd(s, l)
}
fn bwd(&self, arg: &Tensor, _res: &Tensor, grad_res: &Tensor) -> Result<Tensor> {
fn bwd(&self, arg: &Tensor, _res: &Tensor, grad_res: &Tensor) -> Result<Option<Tensor>> {
let alpha = self.0.alpha;
let bwd = arg.custom_op1(EluBackward { alpha })?;
grad_res.mul(&bwd)
Ok(Some(grad_res.mul(&bwd)?))
}
}