docstr: No special modification required for Lightning, define it as you normally would
  in the `nn.Module` in vanilla PyTorch.
docstr_tok: ["No","special","modification","required","for","Lightning",",","define","it","as","you","normally","would","in","the","`","nn.Module","`","in","vanilla","PyTorch","."]
code: def forward(self, input_ids, attention_mask, token_type_ids):
      """
      No special modification required for Lightning, define it as you normally would
      in the `nn.Module` in vanilla PyTorch.
      """
      hidden_states = self.bert_model(
          input_ids=input_ids,
          token_type_ids=token_type_ids,
          attention_mask=attention_mask,
      )
      if isinstance(hidden_states, tuple):
          hidden_states = hidden_states[0]
      logits = self.classifier(hidden_states=hidden_states)
      return logits.float()
  
imports: import os
  from typing import Dict, List, Optional
  import torch
  from omegaconf import DictConfig
  from pytorch_lightning import Trainer
  from nemo.collections.common.losses import CrossEntropyLoss
  from nemo.collections.nlp.data.text_classification import TextClassificationDataset, calc_class_weights
  from nemo.collections.nlp.metrics.classification_report import ClassificationReport
  from nemo.collections.nlp.models.nlp_model import NLPModel
  from nemo.collections.nlp.modules.common import SequenceClassifier
  from nemo.collections.nlp.parts.utils_funcs import tensor2list
  from nemo.core.classes.common import typecheck
  from nemo.core.classes.exportable import Exportable
  from nemo.utils import logging
label_desc: Use image node_id 3 for calling the TextClassificationModel obj's underlying member method code with example usage: obj.forward(input_ids, attention_mask, token_type_ids) and returns: logits