Yunchao_Liu:

        loss_mat = torch.where(is_valid, loss_mat, torch.zeros(
            loss_mat.shape).to(loss_mat.device).to(loss_mat.dtype))



Do you think is_valid is on the same device?
Also, do you know the line number of the error?