The input tensor is a NCHW dim tensor, and I got a mask tensor with size N1HW  indicates  every pixel’s label.(0 or 1).I’d like to group within each sample by this mask and do  some matrix manipulation.such as
0 group and 1 group’s inner product. Is there any good idea about how to group per sample in a batch?