Are you getting the same shape mismatch error if you are using torch.ones_like(grad) instead of gradient_mask?
If not, then I assume the gradient_mask creation is creating the wrong shape.
If yes, could you post a minimal, executable code snippet, which would reproduce this error?