Update chatNT.py

chatNT.py

@@ -721,6 +721,7 @@ class TorchMultiOmicsModel(PreTrainedModel):
             english_token_ids=english_token_ids,
             projected_bio_embeddings=projected_bio_embeddings,
         )
+        logits = logits.to(torch.float32)
 
         outs = {"logits": logits, "projected_bio_embeddings": projected_bio_embeddings}
 
@@ -927,9 +928,6 @@ class TorchGptGroupedQueryAttention(nn.Module):
         attention_weights = nn.functional.softmax(attention_logits, dim=-1)
         attention_weights = attention_weights.to(values.dtype)
 
-        print(f"Attention weights type : ", attention_weights.dtype)
-        print(f"Values type : ", values.dtype)
-        
         values = torch.einsum("bhtT,bThd->bthd", attention_weights, values)
         values = values.contiguous().view(batch_size, seq_len, -1)