Correctness of the MemoryAsContextTransformer

[haohq19]

Hi, I'm confused with the implementation of the MACTransformer module. In the paper of Titan, the MACTransformer uses the output of attention to update the NeuralMemory. However, in this implementation, I think the update of NeuralMemory (mem.forward) is conducted before attention. I'm confused if there is something I misunderstand.