* [Fix] Make autocast compatible with mps

* Enhance unit test

* fix unit test

* clean the code

* fix unit test