fix #170667 (comment) pytorch/torchtitan#2747 (comment)
For DTensor parameters, torch.zeros_like(unsharded_param) returns a DTensor, leading to fsdp.chunk_cat.default: got mixed torch.Tensor and DTensor
Co-authored-by: Irvingwangjr [email protected]
Pull Request resolved: #183040
Approved by: https://github.com/mansiag05, https://github.com/anshul-si
Co-authored-by: Irvingwangjr [email protected]