fixing bug with creating masks

creating masks was done after swapping axes, which gave wrong masks