If you have a workable example, I’d take a look. It doesn’t sound like any issue I’m aware of. (Which are: 1) NaN if you don’t pass zero_infinity=True (available in master / the nightly builds listed under “preview”) and have impossible targets. 2) A bad loss/gradient for all blank targets.)
Best regards
Thomas