4326582bb1e68480ef79a02abbf4bfacc3aadede,layers/losses.py,MSELossMasked,forward,#MSELossMasked#Any#Any#Any#,39

Before Change

        // target_flat: (batch * max_len, dim)
        target_flat = target.view(-1, target.shape[-1])
        // losses_flat: (batch * max_len, dim)
        losses_flat = functional.mse_loss(
            input, target_flat, size_average=False, reduce=False)
        // losses: (batch, max_len, dim)
        losses = losses_flat.view(*target.size())

        // mask: (batch, max_len, 1)

After Change

        mask = mask.expand_as(input)
        loss = functional.mse_loss(
            input * mask, target * mask, reduction="sum")
        loss = loss / mask.sum()
        return loss

Instances