Upload modeling_fuxitranyu.py with huggingface_hub

modeling_fuxitranyu.py

@@ -30,6 +30,7 @@ from transformers.utils import (
 )
 from .configuration_fuxitranyu import FuxiTranyuConfig
 
+
 try:
     from flash_attn import flash_attn_func, flash_attn_varlen_func
     from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input  # noqa
@@ -1250,10 +1251,10 @@ class FuxiTranyuForCausalLM(FuxiTranyuPreTrainedModel):
             logits = torch.cat(logits, dim=-1)
         else:
             logits = self.lm_head(hidden_states)
-        logits = logits.float()
 
         loss = None
         if labels is not None:
+            logits = logits.float()
             # Shift so that tokens < n predict n
             shift_logits = logits[..., :-1, :].contiguous()
             shift_labels = labels[..., 1:].contiguous()