Exporting PyTorch Lightning model to ONNX format not working
am using Jupyter Lab to run. It has pre-installed tf2.3_py3.6 kernel installed in it. It has 2 GPUS in it.
PyTorch Lightning Version (e.g., 1.3.0): '1.4.6'
PyTorch Version (e.g., 1.8): '1.6.0+cu101'
Python version: 3.6
OS (e.g., Linux): system='Linux'
CUDA/cuDNN version: 11.2
How you installed PyTorch (conda, pip, source): pip

I am saving the best model in checkpoint.
I am doing multi-label classification using Hugging face model. After training the model I want to export the model using ONNX format. The input is attention mask, input ids.
Here is the DataModule Class
N_EPOCHS = 30
BATCH_SIZE = 10

class  SRDataModule(pl.LightningDataModule):
    
    def __init__(self, X_train,y_train, X_test,y_test, tokenizer, batch_size=8, max_token_len=512):
        super().__init__()
        self.batch_size = batch_size
        self.train_df = X_train
        self.test_df = X_test
        self.train_lab = y_train
        self.test_lab = y_test
        self.tokenizer = tokenizer
        self.max_token_len = max_token_len

    def setup(self, stage=None):
        self.train_dataset = SRDataset(
          self.train_df,
          self.train_lab,
          self.tokenizer,
          self.max_token_len
        )

        self.test_dataset = SRDataset(
          self.test_df,
          self.test_lab,
          self.tokenizer,
          self.max_token_len
    )

    def train_dataloader(self):
        return DataLoader(
          self.train_dataset,
          batch_size=self.batch_size,
          shuffle=True,
          num_workers=10
        )

    def val_dataloader(self):
        return DataLoader(
          self.test_dataset,
          batch_size=self.batch_size,
          num_workers=10
        )

    def test_dataloader(self):
        return DataLoader(
          self.test_dataset,
          batch_size=self.batch_size,
          num_workers=10
        )

Here is the model class:
class SRTagger(pl.LightningModule):

  def __init__(self, n_classes: int, n_training_steps=None, n_warmup_steps=None):
    super().__init__()
    self.save_hyperparameters()
    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.BCELoss()

  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.sigmoid(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[&quot;input_ids&quot;]
    attention_mask = batch[&quot;attention_mask&quot;]
    labels = batch[&quot;labels&quot;]
    loss, outputs = self(input_ids, attention_mask, labels)
    
    self.log(&quot;train_loss&quot;, loss, prog_bar=True, logger=True)
    return {&quot;loss&quot;: loss, &quot;predictions&quot;: outputs, &quot;labels&quot;: labels}

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

  def test_step(self, batch, batch_idx):
    input_ids = batch[&quot;input_ids&quot;]
    attention_mask = batch[&quot;attention_mask&quot;]
    labels = batch[&quot;labels&quot;]
    loss, outputs = self(input_ids, attention_mask, labels)
    self.log(&quot;test_loss&quot;, 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[&quot;labels&quot;].detach().cpu():
        labels.append(out_labels)
      for out_predictions in output[&quot;predictions&quot;].detach().cpu():
        predictions.append(out_predictions)

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

    for i, name in enumerate(LABEL_COLUMNS):
      class_roc_auc = auroc(predictions[:, i], labels[:, i])
      self.logger.experiment.add_scalar(f&quot;{name}_roc_auc/Train&quot;, class_roc_auc, self.current_epoch)


  def configure_optimizers(self):

    optimizer = optim.RAdam(self.parameters(), lr=2e-4)

    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'
      )
    )

Sample Data
sample_batch = next(iter(DataLoader(train_dataset, batch_size=10, num_workers=2)))
sample_batch[&quot;input_ids&quot;].shape, sample_batch[&quot;attention_mask&quot;].shape

(torch.Size([10, 512]), torch.Size([10, 512]))
sample_batch.keys()
dict_keys(['text_data', 'input_ids', 'attention_mask', 'labels'])

Model
model = SRTagger(
  n_classes=100,
  n_warmup_steps=warmup_steps,
  n_training_steps=total_training_steps 
)

ONNX code
# # Export the model
torch.onnx.export(model,                     # model being run
                  ##since model is in the cuda mode, input also need to be
                  (sample_batch[&quot;input_ids&quot;],sample_batch[&quot;attention_mask&quot;]),              # model input (or a tuple for multiple inputs)
                  &quot;model_torch_export.onnx&quot;, # where to save the model (can be a file or file-like object)
                  export_params=True,        # store the trained parameter weights inside the model file
                  opset_version=10,          # the ONNX version to export the model to
                  do_constant_folding=True,  # whether to execute constant folding for optimization
                  input_names = ['input'],   # the model's input names
                  output_names = ['output'], # the model's output names
                  dynamic_axes={'input' : {0 : 'batch_size'},    # variable lenght axes
                                'output' : {0 : 'batch_size'}})

Error
RuntimeError: output 1 (0
[ CPULongType{} ]) of traced region did not have observable data dependence with trace inputs; this probably indicates your program cannot be understood by the tracer.

Cuz at onnx conversion, the model output must be tensor, not list, tuple, or dict. You can check your forward of model.
