Hi,
Assuming this is not a differentiable output, you can return it by side effect:
    def custom_dec(self, module, out_dict):
        def custom_forward(*inputs):
            output = module(inputs[0], inputs[1],
                            encoder_attn_mask=inputs[2],
                            decoder_padding_mask=inputs[3],
                            layer_state=inputs[4],
                            causal_mask=inputs[5],
                            output_attentions=inputs[6],
                            )
            # output[2] is a python dictionary
            out_dict[0] = output[2]
            return output[0]

out_dict = []
x = checkpoint.checkpoint(
                    self.custom_dec(decoder_layer, out_dict),
                    x,
                    encoder_hidden_states,
                    encoder_padding_mask,
                    decoder_padding_mask,
                    layer_state,
                    decoder_causal_mask,
                    output_attentions,
                )
layer_past = out_dict[0]