Does it support fp16 training?

[SkyAndCloud]

Hi, thank you for the efforts. I want to use fast_soft_sort in fp16 setting. It seems that fast_soft_sort doesn't support fp16 yet? Looking forward to your reply!

[mblondel]

Is it for tensorflow, pytorch or JAX?

CC @josipd

[SkyAndCloud]

Thank you for the quick reply. This question is for pytorch