Solved by ptrblck in post #2 
              
              
                The code looks generally alright. 
Based on the logic you are using, I assume you don’t want to train the latest_tensor? 
If that’s correct, you could register is as a buffer (self.register_buffer('latest_tensor', torch.zeros(...))) and store the detached version of the input: 
self.latest_tensor = …