Merge pull request #47 from lucidrains/fix

fix types

Author: lucidrains

.github/workflows/test.yml

@@ -17,7 +17,7 @@ jobs:
         run: |
           python -m pip install uv
           python -m uv pip install --upgrade pip
-          python -m uv pip install torch --index-url https://download.pytorch.org/whl/nightly/cpu
+          python -m uv pip install torch==2.8.0 --index-url https://download.pytorch.org/whl/cpu
           python -m uv pip install -e .[test]
       - name: Test with pytest
         run: |

transfusion_pytorch/transfusion.py

@@ -2176,7 +2176,9 @@ def forward(
         tuple[tuple[Float['b _ _'], GetPredFlows], Tensor] |
         Scalar |
         tuple[Scalar, LossBreakdown] |
-        list[Float['b _ _']]
+        list[Float['b _ _']] |
+        tuple[Float['b _ l'], Tensor] |
+        list[list[Tensor]] # predicted flows from return_only_pred_flows = True
     ):
         is_decoding = exists(decoding_text_or_modality)