add dtype information for input (#55358)
Summary:
add dtype for all input besides input dimenstion.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/55358
Reviewed By: heitorschueroff
Differential Revision: D27862346
Pulled By: ilia-cher
fbshipit-source-id: 656c5d6c9f23d723b27b44f0afc1a249ce1f3e44
