define the mask function. By using Maks, we do not need to call Multiply to mask tensor and it is cheaper
