Update modules.py

making code suitable for different batch sizes.
icon-lab · Dec 3, 2024 · 12f6739 · 12f6739
1 parent 8500de8
commit 12f6739
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/models/modules.py b/models/modules.py
@@ -473,7 +473,7 @@ def forward(self, x):
 
         concatenated = torch.cat((out1, out2), dim=1)
         output = self.conv1d(concatenated) 
-        mamba_out = output.view(1, 256, 64, 64)
+        mamba_out = output.view(-1, 256, 64, 64)
 
 
         return mamba_out