class_docstr: Implements a ViTSTR architecture as described in `"Vision Transformer for Fast and
  Efficient Scene Text Recognition" <https://arxiv.org/pdf/2105.08582.pdf>`_.
  
  Args:
  ----
      feature_extractor: the backbone serving as feature extractor
      vocab: vocabulary used for encoding
      embedding_units: number of embedding units
      max_length: maximum word length handled by the model
      dropout_prob: dropout probability for the encoder and decoder
      input_shape: input shape of the image
      exportable: onnx exportable returns only logits
      cfg: dictionary containing information about the model
class_docstr_tok: ["Implements","a","ViTSTR","architecture","as","described","in","`","``","Vision","Transformer","for","Fast","and","Efficient","Scene","Text","Recognition","''","<","https",":","\/\/arxiv.org\/pdf\/2105.08582.pdf",">","`","_",".","Args",":","--","--","feature_extractor",":","the","backbone","serving","as","feature","extractor","vocab",":","vocabulary","used","for","encoding","embedding_units",":","number","of","embedding","units","max_length",":","maximum","word","length","handled","by","the","model","dropout_prob",":","dropout","probability","for","the","encoder","and","decoder","input_shape",":","input","shape","of","the","image","exportable",":","onnx","exportable","returns","only","logits","cfg",":","dictionary","containing","information","about","the","model"]
docstr: Compute categorical cross-entropy loss for the model.
  Sequences are masked after the EOS character.
  
  Args:
  ----
      model_output: predicted logits of the model
      gt: the encoded tensor with gt labels
      seq_len: lengths of each gt word inside the batch
  
  Returns:
  -------
      The loss of the model on the batch
docstr_tok: ["Compute","categorical","cross-entropy","loss","for","the","model",".","Sequences","are","masked","after","the","EOS","character",".","Args",":","--","--","model_output",":","predicted","logits","of","the","model","gt",":","the","encoded","tensor","with","gt","labels","seq_len",":","lengths","of","each","gt","word","inside","the","batch","Returns",":","--","--","--","-","The","loss","of","the","model","on","the","batch"]
code: def compute_loss(
      model_output: tf.Tensor,
      gt: tf.Tensor,
      seq_len: List[int],
  ) -> tf.Tensor:
      """Compute categorical cross-entropy loss for the model.
      Sequences are masked after the EOS character.
  
      Args:
      ----
          model_output: predicted logits of the model
          gt: the encoded tensor with gt labels
          seq_len: lengths of each gt word inside the batch
  
      Returns:
      -------
          The loss of the model on the batch
      """
      # Input length : number of steps
      input_len = tf.shape(model_output)[1]
      # Add one for additional <eos> token (sos disappear in shift!)
      seq_len = tf.cast(seq_len, tf.int32) + 1
      # One-hot gt labels
      oh_gt = tf.one_hot(gt, depth=model_output.shape[2])
      # Compute loss: don't forget to shift gt! Otherwise the model learns to output the gt[t-1]!
      # The "masked" first gt char is <sos>.
      cce = tf.nn.softmax_cross_entropy_with_logits(
          oh_gt[:, 1:, :], model_output
      )
      # Compute mask
      mask_values = tf.zeros_like(cce)
      mask_2d = tf.sequence_mask(seq_len, input_len)
      masked_loss = tf.where(mask_2d, cce, mask_values)
      ce_loss = tf.math.divide(
          tf.reduce_sum(masked_loss, axis=1),
          tf.cast(seq_len, model_output.dtype),
      )
  
      return tf.expand_dims(ce_loss, axis=1)
  
file_path: doctr/doctr/models/recognition/vitstr/tensorflow.py
imports: from copy import deepcopy
  from typing import Any, Dict, List, Optional, Tuple
  import tensorflow
  from tensorflow.keras import Model, layers
  from doctr.datasets import VOCABS
  from ...classification import vit_b, vit_s
  from ...utils.tensorflow import _bf16_to_float32, load_pretrained_params
  from .base import _ViTSTR, _ViTSTRPostProcessor
label_desc: Use image node_id 2 for calling the ViTSTR obj's underlying member method code with example usage: obj.compute_loss(model_output, gt, seq_len) and returns: tf