dtype aligns in nll

[jonas-eschle]

Hi, small fix: the dtype of the 1.0 is float32, which fails if the function is jitted. This just ensures that the 1 has the same dtype as the weigths.

[jiangyi15]

The style check is failed. You can run pre-commit run -a to fix it. See here

[jonas-eschle]

Fair, fixed.

Besides, are you aware of/did you consider to use https://pre-commit.ci/? This would solve such problems and automatically run pre-commit on any commit
(I run pre-commit locally, but this was a quick browser edit -> if you register there, it would have automatically taken care of it)

[jiangyi15]

Thanks for you suggestion, I will have a try.