Update modeling_xlm_roberta.py

modeling_xlm_roberta.py

@@ -61,7 +61,7 @@ except ImportError:
 try:
     from flash_attn.losses.cross_entropy import CrossEntropyLoss
 except ImportError:
-    CrossEntropyLoss = None
+    CrossEntropyLoss = torch.nn.CrossEntropyLoss
 
 try:
     from tqdm.autonotebook import trange
@@ -1168,14 +1168,15 @@ class XLMRobertaClassificationHead(nn.Module):
 
     def __init__(self, config):
         super().__init__()
-        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        linear_cls = nn.Linear if not fused_bias_fc else FusedDense
+        self.dense = linear_cls(config.hidden_size, config.hidden_size)
         classifier_dropout = (
             config.classifier_dropout
             if config.classifier_dropout is not None
             else config.hidden_dropout_prob
         )
         self.dropout = nn.Dropout(classifier_dropout)
-        self.out_proj = nn.Linear(config.hidden_size, config.num_labels)
+        self.out_proj = linear_cls(config.hidden_size, config.num_labels)
 
     def forward(self, features, **kwargs):
         x = features[:, 0, :]  # take <s> token (equiv. to [CLS])