Hello,
I am trying to convert T5 from Google model using HuggingFace library to use gradient checkpointing. However, some variables might be None which lead to an error from checkpointing. Is there any solutions to this issue.
Code:

  
      github.com
  
  
    agemagician/transformers/blob/patch-2/src/transformers/modeling_t5.py#L758 1

all_attentions = () if output_attentions else None
position_bias = None
encoder_decoder_position_bias = None

hidden_states = self.dropout(inputs_embeds)

for i, (layer_module, past_key_value) in enumerate(zip(self.block, past_key_values)):
    if output_hidden_states:
        all_hidden_states = all_hidden_states + (hidden_states,)

    if getattr(self.config, "gradient_checkpointing", False):

        def create_custom_forward(module):
            def custom_forward(*inputs):
                #return module(*inputs, output_attentions)
                print("Started the process !!!")
                return module(*inputs)

            return custom_forward

        print("hidden_states")



  
  
    
    
  
  


Error message :

TypeError: CheckpointFunctionBackward.forward: expected Variable (got NoneType) for return value 1