Skip to content

T
tic
Commit 63a3477a by Pariksheet Pinjari Committed by Tianqi Chen
Options

Update softmax.h (#1057)

parent e15aae2b
Showing with 2 additions and 2 deletions
topi/include/topi/nn/softmax.h
......@@ -83,7 +83,7 @@ inline Tensor softmax(const Tensor &x,
});
return tvm::compute(input_shape, [&](const Array<Var> &indices) {
return _normalize(max_elem, expsum, indices);
});
}, name, tag);
}
/*!
......@@ -116,7 +116,7 @@ inline Tensor log_softmax(const Tensor& x,
return tvm::compute(
x->shape, [&](Var i, Var j) {
return x(i, j) - max_elem(i) - tvm::log(expsum(i));
});
}, name, tag);
}
} // namespace nn
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment