def reset_memory(self, batch_size: Optional[int] = None, state: Optional[torch.Tensor] = None):
        """
        **Overview**:
            Reset the memory of GTrXL, which is called at the beginning of each episode.
            Memory is used to save hidden states of past segments.
        """
        # Reset the memory of GTrXL.
        self.memory = Memory(memory_len=self.memory_len, layer_num=self.layer_num, embedding_dim=self.embedding_dim)
        # If batch_size is not None, specify the batch_size when initializing the memory.
        if batch_size is not None:
            self.memory = Memory(self.memory_len, batch_size, self.embedding_dim, self.layer_num)
        # If state is not None, add state into the memory.
        elif state is not None:
            self.memory.init(state)

        **Overview**:
            Reset the memory of GTrXL, which is called at the beginning of each episode.
            Memory is used to save hidden states of past segments.
        
chapter5_time/gtrxl.py