Fix skip connection

decoder_only_t5/modeling.py

@@ -532,9 +532,9 @@ class DecoderOnlyT5Block(modeling_t5.T5Block):
 
         if self.parallel_layers:
             # https://github.com/google/flaxformer/blob/ea17eb012a1d340ddff017b7a534c2162aaec34c/flaxformer/architectures/t5/t5_architecture.py#L534-L578
-            hidden_states = x + ff_output
-            hidden_states *= 2**-0.5
-            hidden_states = hidden_states + self.layer[0].dropout(hidden_states)
+            x = x + ff_output
+            x *= 2**-0.5
+            hidden_states = hidden_states + self.layer[0].dropout(x)
         else:
             hidden_states = ff_layer(x)