not working with mixed precision

[quocpeyrot]

Hello, this awesome callback doesnt seem to support mixed precision (e.g. learn.to_fp16()).

In my local copy, as a workaround, I hacked the code to support it with the following change, but it significantly slow down the runtime (I can't remember exactly, but ~2x IIRC)

the patch/hack is simple, before the call to histc(), which doesnt support half precision, I just added .float().histc(...)

But it would probably be more efficient to improve histc() instead of casting/converting to float() all the time.