Update core.cairo

gizatechxyz · Mar 25, 2024 · a7d2af9 · a7d2af9
1 parent 61e8542
commit a7d2af9
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/operators/nn/core.cairo b/src/operators/nn/core.cairo
@@ -74,7 +74,7 @@ trait NNTrait<T> {
     /// ## Args
     ///
     /// * `tensor`(`@Tensor<T>`) - The input tensor.
-    /// * `axis`(`usize`) - Describes the dimension Softmax will be performed on. Negative value means counting dimensions from the back. Accepted range is [-r, r-1] where r = rank(input).
+    /// * `axis`(`Option<i32>`) - Describes the dimension Softmax will be performed on. Negative value means counting dimensions from the back. Accepted range is [-r, r-1] where r = rank(input).
     ///
     /// ## Returns
     ///