add model

modeling_cxrbert.py

@@ -18,7 +18,7 @@ BERTTupleOutput = Tuple[T, T, T, T, T]
 
 class CXRBertOutput(ModelOutput):
     last_hidden_state: torch.FloatTensor
-    prediction_logits: torch.FloatTensor
+    logits: torch.FloatTensor
     cls_projected_embedding: Optional[torch.FloatTensor] = None
     hidden_states: Optional[Tuple[torch.FloatTensor]] = None
     attentions: Optional[Tuple[torch.FloatTensor]] = None
@@ -97,7 +97,7 @@ class CXRBertModel(BertForMaskedLM):
         if return_dict:
             return CXRBertOutput(
                 last_hidden_state=last_hidden_state,
-                prediction_logits=bert_for_masked_lm_output.logits,
+                logits=bert_for_masked_lm_output.logits,
                 cls_projected_embedding=cls_projected_embedding,
                 hidden_states=bert_for_masked_lm_output.hidden_states if output_hidden_states else None,
                 attentions=bert_for_masked_lm_output.attentions,