ERROR：train sparse model with FDSP2

[Vieeo]

Bug:

How to solve it?

masked_weight = weight * mod._weight_mask
RuntimeError: aten.mul.Tensor: got mixed torch.Tensor and DTensor, need to convert all torch.Tensor to DTensor before calling distributed operators!

system info:

torch 2.9.0+cu12.6
modelopt 0.41

[Vieeo]

If possible，also fix it with FDSP1：
masked_weight = weight * mod._weight_mask
[rank0]: RuntimeError: The size of tensor a (17698976) must match the size of tensor b (3072) at non-singleton dimension 1