Solved by 111248 in post #2 
              
              
                I solve the problem… 
I forgot to add .data  so that the gradient was ignored… 
class Flip(nn.Module):
    def __init__(self, t, inplace=True):
        super(Flip, self).__init__()
        self.t = t
        self.inplace = inplace

    def forward(self, x):
        x = self.flip(x, self.t, self.inpl…