added padding in fusion.py

auto_mixer/modules/fusion.py

@@ -1,6 +1,7 @@
 import math
 
 import torch
+import torch.nn.functional as F
 from torch import nn
 
 
@@ -114,6 +115,7 @@ def __init__(self, dim=1, **_kwargs):
         self.dim = dim
 
     def __call__(self, *args):
+        args = pad_tensors(args)
         return torch.cat(args, dim=self.dim)
 
     def get_output_shape(self, *args, dim=None):
@@ -146,6 +148,15 @@ def get_output_shape(self, *args, dim=None):
         return tuple(shape)
 
 
+def pad_tensors(args):
+    max_size = max(arg.size(2) for arg in args)
+    padded_tensors = []
+    for tensor in args:
+        padding = (0, max_size - tensor.size(2))
+        padded_tensors.append(F.pad(tensor, padding))
+    return padded_tensors
+
+
 class ConcatDynaFusion:
     def __init__(self, dim=1, **_kwargs):
         self.dim = dim