ability to return outputs from autoregressive wrapper forward

lucidrains · Nov 25, 2023 · 8795108 · 8795108
1 parent a00f542
commit 8795108
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'x-transformers',
   packages = find_packages(exclude=['examples']),
-  version = '1.25.8',
+  version = '1.25.9',
   license='MIT',
   description = 'X-Transformers - Pytorch',
   author = 'Phil Wang',

diff --git a/x_transformers/autoregressive_wrapper.py b/x_transformers/autoregressive_wrapper.py
@@ -254,7 +254,7 @@ def generate(
 
         return out
 
-    def forward(self, x, **kwargs):
+    def forward(self, x, return_outputs = False, **kwargs):
         seq, ignore_index, add_attn_z_loss = x.shape[1], self.ignore_index, self.add_attn_z_loss
 
         inp, target = x[:, :-1], x[:, 1:]
@@ -284,4 +284,7 @@ def forward(self, x, **kwargs):
         if add_attn_z_loss:
             loss = loss + cache.attn_z_loss
 
-        return loss
+        if not return_outputs:
+            return loss
+
+        return loss, (logits, cache)