returning logits and kv cache

Author: lucidrains

transfusion_pytorch/transfusion.py

@@ -2176,7 +2176,8 @@ def forward(
         tuple[tuple[Float['b _ _'], GetPredFlows], Tensor] |
         Scalar |
         tuple[Scalar, LossBreakdown] |
-        list[Float['b _ _']]
+        list[Float['b _ _']] |
+        tuple[Float['b _ l'], Tensor]
     ):
         is_decoding = exists(decoding_text_or_modality)