A minor question about the implementation in moba_naive.py

[zhengqigao]

Thanks for the great work. I wonder why we need need_attend = torch.logical_and(need_attend, gate_idx_mask) in moba_naive.py. It seems to me we can directly take gate_idx_mask as the variable need_attend.