Update modernberg_model.py

modernberg_model.py

@@ -634,7 +634,7 @@ class ModernBergAttention(nn.Module):
         hidden_states = attn_outputs[0]
         hidden_states = self.out_drop(self.Wo(hidden_states))
 
-        return (hidden_states,) + attn_outputs[1:]  # add attentions if outputted
+        return hidden_states
 
 class ModernBergTemporalLayer(nn.Module):
     def __init__(self, config: ModernBergConfig, layer_id: Optional[int] = None):