question: Exporting PyTorch Lightning model to ONNX format not working
context: I 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
  
  Here is the screenshot of my model and it got interrupted due to connection issue.
  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.
  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["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)
  
      for i, name in enumerate(LABEL_COLUMNS):
        class_roc_auc = auroc(predictions[:, i], labels[:, i])
        self.logger.experiment.add_scalar(f"{name}_roc_auc/Train", 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["input_ids"].shape, sample_batch["attention_mask"].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["input_ids"],sample_batch["attention_mask"]),              # model input (or a tuple for multiple inputs)
                    "model_torch_export.onnx", # 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.
answer: Hi
  A google search reveals some help on this issue here:
  pytorch/pytorch#31591
  Citing the thread there
  
  As the error message indicates, the tracer detected that the output of your model didn't have any relationship to the input.
  
  If we look closer at your code, we see that loss=0 and labels=None.
    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
  the if condition does not hold, so the part of your output (the loss) cannot be traced back to any inputs by onnx.
  Change your code to something like this and try again please:
      if labels is not None:
          loss = self.criterion(output, labels)
          return loss, output
      return output