def vmap_t_dim(module: Type[nn.Module]) -> Type[nn.Module]:
  """Vmap a 2D model to add the T dimension."""
  return nn.vmap(
      module,
      in_axes=1,
      out_axes=1,
      variable_axes={'params': None},
      split_rngs={'params': False})
Vmap a 2D model to add the T dimension.
videogvt/models/model_utils.py