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 = input.detach().clone()