staturecrane commited on
Commit
e3fa43c
1 Parent(s): 70bb08a

adds _no_split_block

Browse files

adds _no_split_block to EncoderDecoder model for auto device mapping

Files changed (1) hide show
  1. modeling_codet5p.py +2 -1
modeling_codet5p.py CHANGED
@@ -755,7 +755,8 @@ def shift_tokens_right(input_ids: torch.Tensor, pad_token_id: int, decoder_start
755
  # Adapted from transformers.models.encoder_decoder.modeling_encoder_decoder.EncoderDecoderModel
756
  class CodeT5pEncoderDecoderModel(PreTrainedModel):
757
  config_class = CodeT5pConfig
758
-
 
759
  def __init__(
760
  self,
761
  config: Optional[PretrainedConfig] = None,
 
755
  # Adapted from transformers.models.encoder_decoder.modeling_encoder_decoder.EncoderDecoderModel
756
  class CodeT5pEncoderDecoderModel(PreTrainedModel):
757
  config_class = CodeT5pConfig
758
+ _no_split_modules = ["CodeT5pBlock"]
759
+
760
  def __init__(
761
  self,
762
  config: Optional[PretrainedConfig] = None,