Questions about MAE Loss

[ALEX13679173326]

Great works!
However, I also wonder that the reconstruction loss (MSE Loss) in MAE is lost in the code. I couldn't find it.
Does the MDT computes the noising prediction MSE loss twice in training process (one for complete images and the other for masked image)?

[gasvn]

MDT/masked_diffusion/train_util.py

Line 227 in 4b8bcbb

def forward_backward(self, batch, cond):

Check this function.