fix flagai/model/llama_model.py.

Signed-off-by: ftgreat <[email protected]>
silchan · May 26, 2023 · 0206658 · 0206658
1 parent d931619
commit 0206658
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/flagai/model/llama_model.py b/flagai/model/llama_model.py
@@ -219,9 +219,8 @@ def forward(self, input_ids: torch.Tensor, start_pos=0, labels=None, **kwargs):
                 layer.start_pos = start_pos
                 h = layer(h, freqs_cis, mask)
 
-
+        h = self.norm(h)
         if labels is not None:
-            h = self.norm(h)
             if self.config.checkpoint_activations:
                 h = checkpoint(create_custom_forward(self.output),h)
             else:
@@ -254,7 +253,6 @@ def forward(self, input_ids: torch.Tensor, start_pos=0, labels=None, **kwargs):
                 'hidden_states': h,
             }
         else :
-
             output = self.output(h[:, -1, :])  # only compute last logits
             return {
                 "logits": output.float()