Update softmax.cairo

gizatechxyz · Oct 22, 2023 · 5f7e22e · 5f7e22e
1 parent 541f9c4
commit 5f7e22e
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/src/operators/nn/functional/softmax.cairo b/src/operators/nn/functional/softmax.cairo
@@ -53,9 +53,10 @@ use orion::operators::tensor::{
 };
 
 /// Cf: NNTrait::softmax docstring
-fn softmaxWide2(z: @Tensor<FP16x16>, axis: usize) -> Tensor<FP16x16W> {
-    let exp_tensor: Tensor<FP16x16W> = exp_upcast(*z);
+fn softmaxWide2(z: @Tensor<FP16x16>, axis: usize) -> Tensor<FP16x16> {
+    // let exp_tensor: Tensor<FP16x16W> = exp_upcast(*z);
     // let sum = exp_tensor.reduce_sum(axis, true);
     // let softmax = exp_tensor / sum;
-    return exp_tensor;
+    // return exp_tensor;
+    *z
 }