Error converting cmpk to onnx format after fine tuning

Beginners
1

I’m trying to change the pytorch-lightning model, (ckpt) to onnx with the following code:

trained_model = Tagger.load_from_checkpoint(
  trainer.checkpoint_callback.best_model_path,
  n_classes=num_calsess
)
trained_model.eval()
trained_model.freeze()



dummy_model_input = tokenizer("This is a sample", return_tensors="pt")

trained_model.to_onnx( "model_lightnining_export.onnx",
                     tuple(dummy_model_input.values()) , export_params=True ,
                       input_names=['input_ids', 'attention_mask'],
                       output_names=['logits'],
                      dynamic_axes={'input_ids': {0: 'batch_size', 1: 'sequence'}, 
                    'attention_mask': {0: 'batch_size', 1: 'sequence'}, 
                    'logits': {0: 'batch_size', 1: 'sequence'}}, 
                   do_constant_folding=True, 
                   opset_version=13, 
                        )

And get the following error (although I did put the attention_mask as a example input):

/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    725             Module: self
    726         """
--> 727         return self._apply(lambda t: t.xpu(device))
    728 
    729     def cpu(self: T) -> T:

TypeError: forward() missing 1 required positional argument: 'attention_mask'

The class of classifier:

class Tagger(pl.LightningModule):

  def __init__(self, n_classes: int, n_training_steps=None, n_warmup_steps=None):
    super().__init__()
    self.bert = BertModel.from_pretrained(BERT_MODEL_NAME, return_dict=True)
    self.classifier = nn.Linear(self.bert.config.hidden_size, n_classes)
    self.n_training_steps = n_training_steps
    self.n_warmup_steps = n_warmup_steps
    self.criterion = nn.CrossEntropyLoss(reduction="mean")
    self.n_classes = n_classes
    self.activation = nn.Softmax(dim=1)

  def forward(self, input_ids, attention_mask, labels=None):
    output = self.bert(input_ids, attention_mask=attention_mask)
    output = self.classifier(output.pooler_output)
    output = torch.relu(output)
    loss = 0
    if labels is not None:
        loss = self.criterion(output, labels)
    return loss, output

  def training_step(self, batch, batch_idx):
    input_ids = batch["input_ids"]
    attention_mask = batch["attention_mask"]
    labels = batch["labels"]
    loss, outputs = self(input_ids, attention_mask, labels)
    self.log("train_loss", loss, prog_bar=True, logger=True)
    return {"loss": loss, "predictions": outputs, "labels": labels}

  def validation_step(self, batch, batch_idx):
    input_ids = batch["input_ids"]
    attention_mask = batch["attention_mask"]
    labels = batch["labels"]
    loss, outputs = self(input_ids, attention_mask, labels)
    self.log("val_loss", loss, prog_bar=True, logger=True)
    return loss

  def test_step(self, batch, batch_idx):
    input_ids = batch["input_ids"]
    attention_mask = batch["attention_mask"]
    labels = batch["labels"]
    loss, outputs = self(input_ids, attention_mask, labels)
    self.log("test_loss", loss, prog_bar=True, logger=True)
    return loss

  def training_epoch_end(self, outputs):
    
    labels = []
    predictions = []
    for output in outputs:
      for out_labels in output["labels"].detach().cpu():
        labels.append(out_labels)
      for out_predictions in output["predictions"].detach().cpu():
        predictions.append(out_predictions)

    labels = torch.stack(labels).int()
    predictions = torch.stack(predictions)

  def configure_optimizers(self):

    optimizer = AdamW(self.parameters(), lr=2e-5)

    scheduler = get_linear_schedule_with_warmup(
      optimizer,
      num_warmup_steps=self.n_warmup_steps,
      num_training_steps=self.n_training_steps
    )

    return dict(
      optimizer=optimizer,
      lr_scheduler=dict(
        scheduler=scheduler,
        interval='step'
      )
    )