diff --git a/switch/module.py b/switch/module.py
index ec42d30..fc58dde 100644
--- a/switch/module.py
+++ b/switch/module.py
@@ -75,7 +75,7 @@ class SwitchFeedForward(nn.Module):
         ]
 
         # Initialize an empty tensor to store outputs
-        final_output = x.new_zeros((batch_size, seq_len, self.dim_feedforward))
+        final_output = x.new_zeros((batch_size * seq_len, self.dim_feedforward))
 
         # Capacity of each expert.
         # $$\mathrm{expert\;capacity} =