Update modeling_xlm_roberta.py

modeling_xlm_roberta.py

@@ -210,12 +210,10 @@ class XLMRobertaEncoder(nn.Module):
         subset_mask: (batch, seqlen), dtype=torch.bool
         """
         if key_padding_mask is None or not self.use_flash_attn:
-            mixer_kwargs = (
-                {"key_padding_mask": key_padding_mask.bool()}
-                if key_padding_mask is not None
-                else None
-            )
-            mixer_kwargs['task_type'] = task_type
+            mixer_kwargs = {'task_type': task_type}
+            if key_padding_mask is not None:
+                mixer_kwargs['key_padding_mask'] = key_padding_mask.bool()
+            
             for layer in self.layers:
                 if self._grad_checkpointing:
                     hidden_states = torch.utils.checkpoint.checkpoint(