value dim defined as dimension out // heads

Author: lucidrains

README.md

@@ -20,7 +20,6 @@ layer = LambdaLayer(
     n = 64 * 64,    # number of input pixels (64 x 64 image)
     m = 64 * 64,    # number of context (64 x 64 global)
     dim_k = 16,     # key dimension
-    dim_v = 64,     # value dimension
     heads = 4,      # number of heads, for multi-query
     dim_u = 1       # 'intra-depth' dimension
 )

lambda_networks/lambda_networks.py

@@ -21,17 +21,17 @@ def __init__(
         n,
         m,
         dim_k,
-        dim_v,
         heads = 4,
         dim_out = None,
         dim_u = 1):
         super().__init__()
-        assert (dim_v % heads) == 0, 'values dimension must be divisible by number of heads for multi-head query'
         dim_out = default(dim_out, dim)
         self.u = dim_u # intra-depth dimension
         self.heads = heads
 
-        dim_v //= heads
+        assert (dim_out % heads) == 0, 'values dimension must be divisible by number of heads for multi-head query'
+        dim_v = dim_out // heads
+
         self.to_q = nn.Conv2d(dim, dim_k * heads, 1, bias = False)
         self.to_k = nn.Conv2d(dim, dim_k * dim_u, 1, bias = False)
         self.to_v = nn.Conv2d(dim, dim_v * dim_u, 1, bias = False)

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'lambda-networks',
   packages = find_packages(),
-  version = '0.0.1',
+  version = '0.0.2',
   license='MIT',
   description = 'Lambda Networks - Pytorch',
   author = 'Phil Wang',