self.save_for_backward (or ctx.save_for_backward in the new function API) is needed for saving any tensors that are one of the inputs to the function; for intermediate values (including noise masks) you should use assignment to self.attr or ctx.attr instead.