by default, dimension of memories (values) will be chosen so output dimensions is equal to the embedding dimension

lucidrains · lucidrains · commit 9d597af6cd85 · 2022-07-25T15:29:38.000-07:00
diff --git a/discrete_key_value_bottleneck_pytorch/discrete_key_value_bottleneck.py b/discrete_key_value_bottleneck_pytorch/discrete_key_value_bottleneck.py
@@ -36,7 +36,7 @@ def __init__(
             **kwargs
         )
 
-        dim_memory = default(dim_memory, dim)
+        dim_memory = default(dim_memory, dim // num_memory_codebooks)
         self.values = nn.Parameter(torch.randn(num_memory_codebooks, num_memories, dim_memory))
 
     def forward(self, x, **kwargs):
diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'discrete-key-value-bottleneck-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.0.3',
+  version = '0.0.4',
   license='MIT',
   description = 'Discrete Key / Value Bottleneck - Pytorch',
   author = 'Phil Wang',

Original file line number	Diff line number	Diff line change
`@@ -36,7 +36,7 @@ def __init__(`
`36`	`36`	`**kwargs`
`37`	`37`	`)`
`38`	`38`
`39`		`- dim_memory = default(dim_memory, dim)`
	`39`	`+ dim_memory = default(dim_memory, dim // num_memory_codebooks)`
`40`	`40`	`self.values = nn.Parameter(torch.randn(num_memory_codebooks, num_memories, dim_memory))`
`41`	`41`
`42`	`42`	`def forward(self, x, **kwargs):`