Solved by ptrblck in post #2 
              
              
                Try to call cache = copy.deepcopy(self.state_dict()), as currently cache holds a reference to your state dict and will be updated with it.