code: def forward(
      self,
      text_tokens: Tensor,
      key_padding_mask: Optional[Tensor] = None,
      *args,
      **kwargs,
  ) -> Tensor:
      if text_tokens.dim() == 4:
          # It's for zero-shot evaluation.
          # Each class in the dataset has multiple captions
          # Encoding happens separately for each classes/captions due to OOM issue
          return self.forward_zero_shot(
              text_tokens=text_tokens,
              key_padding_mask=key_padding_mask,
              *args,
              **kwargs,
          )
      elif text_tokens.dim() == 2:
          # Image-text pair data with single caption
          # [B, CL] --> [B, d]
          text_tokens = self.encode_text(
              text_tokens=text_tokens,
              key_padding_mask=key_padding_mask,
              *args,
              **kwargs,
          )
          return text_tokens
      elif text_tokens.dim() == 3:
          # Image-text pair with multiple captions per image (e.g. Flickr-30k)
          # Treat them as separate captions by reshaping into batch dim
          # [B, N, C] --> [B*N, C] -encode-> [B*N, d] --> [B, N, d]
          b, n, _ = text_tokens.shape
          text_tokens = text_tokens.reshape(b * n, -1)
          if key_padding_mask:
              key_padding_mask = key_padding_mask.reshape(b * n, -1)
          text_tokens = self.encode_text(
              text_tokens=text_tokens,
              key_padding_mask=key_padding_mask,
              *args,
              **kwargs,
          )
          text_tokens = text_tokens.reshape(b, n, -1)
          return text_tokens
      else:
          raise NotImplementedError
  
imports: import argparse
  import math
  from typing import Optional, Sequence
  import torch
  from torch import Tensor, nn
  from torch.nn import functional
  from torch.utils.checkpoint import checkpoint
  from cvnets.layers import Dropout, Embedding, PositionalEmbedding, get_normalization_layer
  from cvnets.modules import TransformerEncoder
  from cvnets.text_encoders import TEXT_ENCODER_REGISTRY, BaseTextEncoder
  from utils import logger
label_desc: Use image node_id 8 for calling the TextTransformer obj's underlying member method code with example usage: obj.forward(text_tokens, key_padding_mask) and returns: self, text_tokens, text_tokens