examples/research_projects/onnx/summarization/bart_onnx/generation_onnx.py/BARTGenerator/_trace_modules
class BARTGenerator:    def _trace_modules(self, model):
        input_ids = torch.tensor(
            [
                [
                    19,
                    669,
                    18,
                    420,
                    8,
                    664,
                    57,
                    42,
                    8,
                    664,
                    21,
                    3028,
                    195,
                    4445,
                    331,
                    1293,
                    34,
                    21,
                    10,
                    6174,
                    1100,
                    6,
                    69,
                    104,
                    42,
                    32,
                    2621,
                    1638,
                    144,
                    4,
                    6174,
                    558,
                    108,
                    4419,
                    1091,
                    28,
                    4,
                    1668,
                    9,
                    1509,
                    1621,
                    279,
                    35,
                    867,
                    2734,
                    85,
                    11,
                    2216,
                    2734,
                    85,
                    203,
                    2244,
                    7,
                    6,
                    15,
                    8102,
                    7,
                    57,
                    8629,
                    5,
                    model.config.eos_token_id,
                ]
            ],
            device=model.device,
            dtype=torch.long,
        )
        attention_mask = torch.tensor(
            [[True] * input_ids.shape[-1]],
            device=model.device,
            dtype=torch.bool,
        )
        self.encoder = _create_traced_encoder(model.get_encoder(), input_ids, attention_mask)
        encoder_outputs = model.get_encoder()(input_ids, attention_mask=attention_mask, return_dict=True)
        decoder = model.model.decoder
        decoder_outputs = decoder(input_ids, attention_mask, encoder_outputs["last_hidden_state"], None, None, None)
        self.decoder_no_past = _create_traced_decoder(
            model.model.decoder, input_ids, encoder_outputs["last_hidden_state"], attention_mask
        )
        self.decoder_with_past = _create_traced_decoder(
            model.model.decoder, input_ids, encoder_outputs["last_hidden_state"], attention_mask, decoder_outputs[1]
        )