Improve torch_dtype checks (#40808)
* Improve torch_dtype checks
Signed-off-by: Yuanyuan Chen <cyyever@outlook.com>
* Apply suggestions from code review
---------
Signed-off-by: Yuanyuan Chen <cyyever@outlook.com>
Co-authored-by: Joao Gante <joaofranciscocardosogante@gmail.com>
