code: def forward(
      self, features: torch.Tensor
  ) -> Tuple[torch.Tensor, torch.Tensor, List[torch.Tensor]]:
      if features.dim() != 2:
          raise ValueError(
              f"Expecting incoming tensor to be dim 2, "
              f"instead dim={features.dim()}"
          )
  
      # shape notation
      # i_s: input_size
      # s: size
      # o_s: output_size
      # b_s: batch_size
      batch_size = features.shape[0]  # b_s
      # Tile out_accumulator, aggregated_mask, and prior_scales to add batch dimension.
      out_accumulator = torch.tile(
          self.out_accumulator, (batch_size, 1)
      )
      aggregated_mask = torch.tile(
          self.aggregated_mask, (batch_size, 1)
      )
      prior_scales = torch.tile(self.prior_scales, (batch_size, 1))
      masks = []
      total_entropy = 0.0
  
      if batch_size != 1 or not self.training:
          # Skip batch normalization training if the batch size is 1.
          features = self.batch_norm(features)  # [b_s, i_s]
      elif batch_size == 1:
          # We temporarily set the batch_norm module to eval mode as we can't compute the running statistics
          # when the batch size is 1.
          self.batch_norm.eval()
          features = self.batch_norm(features)  # [b_s, i_s]
          self.batch_norm.train()
      masked_features = features
  
      x = self.feature_transforms[0](masked_features)  # [b_s, s + o_s]
  
      for step_i in range(1, self.num_steps + 1):
          #########################
          # Attentive Transformer #
          #########################
          # x in following is shape [b_s, s]
          mask_values = self.attentive_transforms[step_i](
              x[:, self.output_size :], prior_scales
          )  # [b_s, i_s]
  
          # relaxation factor 1 forces the feature to be only used once
          prior_scales = prior_scales * (
              self.relaxation_factor - mask_values
          )  # [b_s, i_s]
  
          # entropy is used to penalize the amount of sparsity
          # in feature selection
          if self.sparsity.item() != 0.0:
              total_entropy += (
                  torch.mean(
                      torch.sum(
                          -mask_values
                          * torch.log(mask_values + 0.00001),
                          dim=1,
                      )
                  )
                  / self.num_steps
              )
  
          masks.append(
              torch.unsqueeze(torch.unsqueeze(mask_values, 0), 3)
          )  # [1, b_s, i_s, 1]
  
          #######################
          # Feature Transformer #
          #######################
          masked_features = torch.multiply(mask_values, features)
  
          x = self.feature_transforms[step_i](
              masked_features
          )  # [b_s, s + o_s]
  
          # x in following is shape [b_s, o_s]
          out = nn.functional.relu(
              x[:, : self.output_size]
          )  # [b_s, o_s]
          out_accumulator += out
  
          # Aggregated masks are used for visualization of the
          # feature importance attributes.
          scale = torch.sum(out, dim=1, keepdim=True) / self.num_steps
          aggregated_mask += mask_values * scale  # [b_s, i_s]
  
      final_output = self.final_projection(
          out_accumulator
      )  # [b_s, o_s]
  
      sparsity_loss = torch.multiply(self.sparsity, total_entropy)
      self.update_loss("sparsity_loss", sparsity_loss)
  
      return final_output, aggregated_mask, masks
  
imports: from typing import List, Optional, Tuple
  import torch
  import torch.nn
  from ludwig.modules.normalization_modules import GhostBatchNormalization
  from ludwig.utils.entmax import Entmax15, EntmaxBisect, Sparsemax
  from ludwig.utils.torch_utils import LudwigModule
label_desc: Use image node_id 2 for calling the TabNet obj's underlying member method code with example usage: obj.forward(features) and returns: final_output,  aggregated_mask,  masks