Prefer not to use for loops, try to vectorize your code as much as possible.
Refer this old question I asked for checking my implementation of softargmax which I believe you are lookin for. It’s pretty decent and fast…
1 Like