* Add casting initializers for bf16 training

* Changes after review

* Correct test and add comment
Co-authored-by: Njoanna.wozna.intel <joanna.wozna@intel.com>