这个注意力码部分看起来不正确

attention_mlp = basic.BatchApply(mlp.MLP([self._mem_size] * self._attention_mlp_layers))

for _ in range(self._num_blocks): attended_memory = self._multihead_attention(memory)

应该是这个吧