Skip to content

Commit

Permalink
reinject the write memory positions
Browse files Browse the repository at this point in the history
  • Loading branch information
lucidrains committed Aug 29, 2023
1 parent 8b82aaa commit 8f9c3ad
Show file tree
Hide file tree
Showing 2 changed files with 4 additions and 4 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -337,10 +337,10 @@ def forward(

# prepare write memories, as in paper

write_memories = self.init_memory(b)

if exists(read_memories) and self.add_write_to_next_write_mem:
write_memories = read_memories
else:
write_memories = self.init_memory(b)
write_memories = write_memories + read_memories

# prepare read memories

Expand Down
2 changes: 1 addition & 1 deletion setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
setup(
name = 'recurrent-memory-transformer-pytorch',
packages = find_packages(exclude=[]),
version = '0.5.1',
version = '0.5.2',
license='MIT',
description = 'Recurrent Memory Transformer - Pytorch',
author = 'Phil Wang',
Expand Down

0 comments on commit 8f9c3ad

Please sign in to comment.