gmastrapas commited on
Commit
bcfceeb
1 Parent(s): 12700ba

Add `add_pooling_layer` arg to XLMRobertaLora

Browse files
Files changed (1) hide show
  1. modeling_lora.py +5 -2
modeling_lora.py CHANGED
@@ -250,11 +250,14 @@ class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
250
  """
251
 
252
  def __init__(
253
- self, config: XLMRobertaFlashConfig, roberta: Optional[XLMRobertaModel] = None
 
 
 
254
  ):
255
  super().__init__(config)
256
  if roberta is None:
257
- self.roberta = XLMRobertaModel(config)
258
  else:
259
  self.roberta = roberta
260
 
 
250
  """
251
 
252
  def __init__(
253
+ self,
254
+ config: XLMRobertaFlashConfig,
255
+ roberta: Optional[XLMRobertaModel] = None,
256
+ add_pooling_layer: bool = True,
257
  ):
258
  super().__init__(config)
259
  if roberta is None:
260
+ self.roberta = XLMRobertaModel(config, add_pooling_layer=add_pooling_layer)
261
  else:
262
  self.roberta = roberta
263