search

software-mansion-labs / ex_vision

Apache License 2.0
61 stars 5 forks source link

Replace softmax with stable version + small corrections #8

Closed msluszniak closed 4 months ago

msluszniak commented 4 months ago

I propose some minor changes:

  1. To find number of dimensions it is better to use Nx.rank
  2. There is a simplified syntax for creating tensors of given type in Nx. Currently the type is not provided but I think it's better here to be explicit
  3. In utils module there is an unstable version of softmax, I suggest to use the stable one from Axon
msluszniak commented 4 months ago

@mkopcins I added the changes that I suggested in the previous PRs. I'm not sure if the names of the functions are in point, so feel free to suggest new ones.