You do hit PyTorch operations which don’t support bfloat16 on the selected backend yet, so it’s not as easy.
There are two parts to this:

Personally, I tend to try to look into fixing stuff that I need and doesn’t work. If the backend libraries (eg cublas) support it, it might be doable (I think a relatively new cuda/cublas/cudnn version might help).
The other option could be to manually cast only the supported parts and cast the activations in between. This is similar to the automatic multi-precision (AMP) support but for support reasons and less automatic.

Best regards
Thomas