Yeah, I see that the batch dimension seems to be alright in your first first output ([38, 61] ...).
However, why is the mask size [75, 61]? If the batch chunks have a batch size of 38 or 36, I’m not sure why your masks have suddenly a non-matching shape.