diff --git a/perceiver_pytorch/perceiver_io.py b/perceiver_pytorch/perceiver_io.py
index 4694f7b..f5057a2 100644
--- a/perceiver_pytorch/perceiver_io.py
+++ b/perceiver_pytorch/perceiver_io.py
@@ -178,9 +178,6 @@ def forward(
         
         latents = self.decoder_cross_attn(queries, context = x)
 
-        if not exists(latents):
-            return latents
-
         # final linear out
 
         return self.to_logits(latents)
diff --git a/setup.py b/setup.py
index 081608c..7b36dd5 100644
--- a/setup.py
+++ b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'perceiver-pytorch',
   packages = find_packages(),
-  version = '0.5.0',
+  version = '0.5.1',
   license='MIT',
   description = 'Perceiver - Pytorch',
   author = 'Phil Wang',