diff --git a/perceiver_ar_pytorch/autoregressive_wrapper.py b/perceiver_ar_pytorch/autoregressive_wrapper.py
index 107d9c0..41f02fe 100644
--- a/perceiver_ar_pytorch/autoregressive_wrapper.py
+++ b/perceiver_ar_pytorch/autoregressive_wrapper.py
@@ -56,16 +56,16 @@ def generate(
         out = start_tokens
 
         for _ in range(seq_len):
-            out = out[:, -self.max_seq_len:]
-            logits = self.net(out, **kwargs)[:, -1, :]
+            logits = self.net(
+                out[:, -self.max_seq_len:],
+                **kwargs
+            )[:, -1]
 
             filtered_logits = top_k(logits, thres = filter_thres)
             probs = F.softmax(filtered_logits / temperature, dim=-1)
 
             sample = torch.multinomial(probs, 1)
-
             out = torch.cat((out, sample), dim=-1)
-            offset = max(0, offset - 1)
 
             if exists(eos_token):
                 is_eos_token = out == eos_token
diff --git a/setup.py b/setup.py
index 1433117..6153e36 100644
--- a/setup.py
+++ b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'perceiver-ar-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.0.7',
+  version = '0.0.8',
   license='MIT',
   description = 'Perceiver AR',
   author = 'Phil Wang',