clean up conditional order

jamestiotio · Dec 27, 2022 · 27e56e3 · 27e56e3
1 parent 03458c5
commit 27e56e3
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/megatron/model/transformer.py b/megatron/model/transformer.py
@@ -666,11 +666,11 @@ def forward(self, x, attention_mask, layer_past=None):
 
  residual = x
  # applies the correct normalization depending on if the norms are tied
- if not self.gpt_j_tied:
- x1, x2 = self.input_layernorm(x), self.post_attention_layernorm(x)
- else:
+ if self.gpt_j_tied:
  x = self.input_layernorm(x)
  x1, x2 = x, x
+ else:
+ x1, x2 = self.input_layernorm(x), self.post_attention_layernorm(x)
 
  # attention operator
  attention_output, attention_bias = self.attention(