staturecrane commited on
Commit
a9b9968
1 Parent(s): a191f78

adds no_split_modules

Browse files

Adds _no_split_modules attribute on the EncoderDecoder class for accelerate support

Files changed (1) hide show
  1. modeling_codet5p.py +1 -0
modeling_codet5p.py CHANGED
@@ -758,6 +758,7 @@ def shift_tokens_right(input_ids: torch.Tensor, pad_token_id: int, decoder_start
758
  # Adapted from transformers.models.encoder_decoder.modeling_encoder_decoder.EncoderDecoderModel
759
  class CodeT5pEncoderDecoderModel(PreTrainedModel):
760
  config_class = CodeT5pConfig
 
761
 
762
  def __init__(
763
  self,
 
758
  # Adapted from transformers.models.encoder_decoder.modeling_encoder_decoder.EncoderDecoderModel
759
  class CodeT5pEncoderDecoderModel(PreTrainedModel):
760
  config_class = CodeT5pConfig
761
+ _no_split_modules = ["CodeT5pBlock"]
762
 
763
  def __init__(
764
  self,