one more fix

Author: lucidrains

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "transfusion-pytorch"
-version = "0.4.11"
+version = "0.4.12"
 description = "Transfusion in Pytorch"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

transfusion_pytorch/transfusion.py

@@ -2370,7 +2370,7 @@ def forward(
 
             velocity_match_losses = []
 
-            for ema_pred_flow, pred_flow, is_one_modality in zip(ema_pred_flows, pred_flows, is_modalities.unbind(dim = 1)):
+            for mod, ema_pred_flow, pred_flow, is_one_modality in zip(self.get_all_modality_info(), ema_pred_flows, pred_flows, is_modalities.unbind(dim = 1)):
 
                 velocity_match_loss = F.mse_loss(
                     pred_flow,
@@ -2380,6 +2380,9 @@ def forward(
 
                 is_one_modality = reduce(is_one_modality, 'b m n -> b n', 'any')
 
+                if mod.channel_first_latent:
+                    velocity_match_loss = rearrange(velocity_match_loss, 'b d ... -> b ... d')
+
                 velocity_match_loss = velocity_match_loss[is_one_modality].mean()
 
                 velocity_match_losses.append(velocity_match_loss)