make it so perceiver io decoding queries can omit the batch dimension

lucidrains · Oct 10, 2021 · 2a1b039 · 2a1b039
1 parent 483a6ed
commit 2a1b039
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 2 deletions.
diff --git a/README.md b/README.md
@@ -65,7 +65,7 @@ model = PerceiverIO(
 )
 
 seq = torch.randn(1, 512, 32)
-queries = torch.randn(1, 128, 32)
+queries = torch.randn(128, 32)
 
 logits = model(seq, queries = queries) # (1, 128, 100) - (batch, decoder seq, logits dim)
 ```

diff --git a/perceiver_pytorch/perceiver_io.py b/perceiver_pytorch/perceiver_io.py
@@ -171,6 +171,11 @@ def forward(
         if not exists(queries):
             return x
 
+        # make sure queries contains batch dimension
+
+        if queries.ndim == 2:
+            queries = repeat(queries, 'n d -> b n d', b = b)
+
         # cross attend from decoder queries to latents
 
         latents = self.decoder_cross_attn(queries, context = x)

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'perceiver-pytorch',
   packages = find_packages(),
-  version = '0.7.4',
+  version = '0.7.5',
   license='MIT',
   description = 'Perceiver - Pytorch',
   author = 'Phil Wang',