fix torch_dtype in estimate memory #5421

Job	Run time
run-tests (latest, test_prod)	1m 47s
run-tests (latest, test_core)	2m 54s
run-tests (latest, test_cli)	2m 35s
run-tests (latest, test_big_modeling)	2m 24s
run-tests (latest, test_deepspeed)	2m 18s
run-tests (latest, test_fsdp)	2m 36s
run-tests (latest, test_example_differences)	2m 49s
run-tests (latest, test_checkpoint_step)	3m 48s
run-tests (latest, test_checkpoint_epoch)	3m 38s
run-tests (latest, test_rest)	9m 51s
run-tests (minimum, test_prod)	2m 44s
run-tests (minimum, test_core)	4m 1s
run-tests (minimum, test_cli)	3m 29s
run-tests (minimum, test_big_modeling)	3m 32s
run-tests (minimum, test_deepspeed)	3m 55s
run-tests (minimum, test_fsdp)	3m 16s
run-tests (minimum, test_example_differences)	3m 12s
run-tests (minimum, test_checkpoint_step)	5m 0s
run-tests (minimum, test_checkpoint_epoch)	4m 50s
run-tests (minimum, test_rest)	12m 15s
	1h 20m 54s

Provide feedback