SYMBOL INDEX (305 symbols across 37 files)

FILE: examples/pretrain/preprocess.py
  class PreprocessingConfig (line 32) | class PreprocessingConfig:
  function preprocess (line 55) | def preprocess(cfg: PreprocessingConfig) -> None:

FILE: examples/pretrain/pretrain.py
  class PretrainConfig (line 49) | class PretrainConfig:
  function pretrain (line 79) | def pretrain(cfg: PretrainConfig) -> None:

FILE: examples/usage.py
  function usage (line 17) | def usage() -> None:

FILE: examples/verification/verify.py
  function verify (line 24) | def verify() -> None:

FILE: examples/xla-reference/xpreprocess.py
  class PreprocessingConfig (line 31) | class PreprocessingConfig:
  function xpreprocess (line 54) | def xpreprocess(cfg: PreprocessingConfig) -> None:

FILE: examples/xla-reference/xpretrain.py
  class PretrainConfig (line 65) | class PretrainConfig:
  function xpretrain (line 97) | def xpretrain(cfg: PretrainConfig) -> None:
  function mp_fn (line 819) | def mp_fn(_: int, cfg: PretrainConfig) -> None:
  function main (line 827) | def main(cfg: PretrainConfig) -> None:

FILE: voltron/conf/accelerators.py
  class AcceleratorConfig (line 16) | class AcceleratorConfig:
  class TPUv2OneConfig (line 23) | class TPUv2OneConfig(AcceleratorConfig):
  class TPUv2EightConfig (line 30) | class TPUv2EightConfig(AcceleratorConfig):
  class TPUv3OneConfig (line 37) | class TPUv3OneConfig(AcceleratorConfig):
  class TPUv3EightConfig (line 44) | class TPUv3EightConfig(AcceleratorConfig):
  class TorchRunDefaultConfig (line 55) | class TorchRunDefaultConfig(AcceleratorConfig):

FILE: voltron/conf/datasets.py
  class DatasetConfig (line 16) | class DatasetConfig:
  class SthSthv2Config (line 61) | class SthSthv2Config(DatasetConfig):

FILE: voltron/conf/models.py
  class ModelConfig (line 16) | class ModelConfig:
  class MVPConfig (line 40) | class MVPConfig(ModelConfig):
  class MVPSmallConfig (line 75) | class MVPSmallConfig(MVPConfig):
  class R3MConfig (line 97) | class R3MConfig(ModelConfig):
  class R3MSmallConfig (line 136) | class R3MSmallConfig(R3MConfig):
  class ResNet3MConfig (line 153) | class ResNet3MConfig(ModelConfig):
  class RN3M50Config (line 187) | class RN3M50Config(ResNet3MConfig):
  class VCondConfig (line 203) | class VCondConfig(ModelConfig):
  class VCondSmallConfig (line 247) | class VCondSmallConfig(VCondConfig):
  class VCondBaseConfig (line 272) | class VCondBaseConfig(VCondConfig):
  class VDualConfig (line 299) | class VDualConfig(ModelConfig):
  class VDualSmallConfig (line 344) | class VDualSmallConfig(VDualConfig):
  class VDualBaseConfig (line 369) | class VDualBaseConfig(VDualConfig):
  class VGenConfig (line 396) | class VGenConfig(ModelConfig):
  class VGen50SmallConfig (line 442) | class VGen50SmallConfig(VGenConfig):
  class VGen50BaseConfig (line 467) | class VGen50BaseConfig(VGenConfig):

FILE: voltron/conf/tracking.py
  class TrackingConfig (line 15) | class TrackingConfig:
  class VoltronTrackingConfig (line 39) | class VoltronTrackingConfig(TrackingConfig):

FILE: voltron/datasets/datasets.py
  class PretrainDataset (line 28) | class PretrainDataset(Dataset):
    method __init__ (line 29) | def __init__(self) -> None:
    method hydrate (line 34) | def hydrate(self, path: Path) -> None:
    method set_epoch (line 43) | def set_epoch(self, epoch: int) -> None:
    method __getitem__ (line 46) | def __getitem__(self, idx: int) -> Tuple[torch.Tensor, ...]:
    method __len__ (line 49) | def __len__(self) -> int:
  class StateDataset (line 53) | class StateDataset(PretrainDataset):
    method __init__ (line 54) | def __init__(self, epoch: int, index_path: Path, img_transform: Compos...
    method set_epoch (line 62) | def set_epoch(self, epoch: int) -> None:
    method __getitem__ (line 78) | def __getitem__(self, idx: int) -> torch.Tensor:
    method __len__ (line 85) | def __len__(self) -> int:
  class StateLanguageDataset (line 89) | class StateLanguageDataset(PretrainDataset):
    method __init__ (line 90) | def __init__(
    method set_epoch (line 109) | def set_epoch(self, epoch: int) -> None:
    method __getitem__ (line 129) | def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor, t...
    method __len__ (line 148) | def __len__(self) -> int:
  class StateOKDataset (line 152) | class StateOKDataset(PretrainDataset):
    method __init__ (line 153) | def __init__(
    method set_epoch (line 172) | def set_epoch(self, epoch: int) -> None:
    method __getitem__ (line 192) | def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor, t...
    method __len__ (line 212) | def __len__(self) -> int:
  class GenStateOKDataset (line 216) | class GenStateOKDataset(PretrainDataset):
    method __init__ (line 217) | def __init__(
    method set_epoch (line 236) | def set_epoch(self, epoch: int) -> None:
    method __getitem__ (line 256) | def __getitem__(
    method __len__ (line 287) | def __len__(self) -> int:
  class QuintetDataset (line 291) | class QuintetDataset(PretrainDataset):
    method __init__ (line 292) | def __init__(self, epoch: int, index_path: Path, img_transform: Compos...
    method set_epoch (line 303) | def set_epoch(self, epoch: int) -> None:
    method __getitem__ (line 319) | def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor, t...
    method __len__ (line 333) | def __len__(self) -> int:
  function get_datasets (line 337) | def get_datasets(

FILE: voltron/datasets/v1/stream_datasets.py
  class PretrainDataset (line 54) | class PretrainDataset(Dataset):
    method set_epoch (line 55) | def set_epoch(self, epoch: int) -> None:
  class StateDataset (line 59) | class StateDataset(PretrainDataset):
    method __init__ (line 60) | def __init__(
    method set_epoch (line 80) | def set_epoch(self, epoch: int) -> None:
    method __getitem__ (line 115) | def __getitem__(self, index: int) -> torch.Tensor:
    method __len__ (line 163) | def __len__(self) -> int:
  class StateLanguageDataset (line 167) | class StateLanguageDataset(PretrainDataset):
    method __init__ (line 168) | def __init__(
    method set_epoch (line 200) | def set_epoch(self, epoch: int) -> None:
    method __getitem__ (line 253) | def __getitem__(self, index: int) -> Tuple[torch.Tensor, torch.Tensor,...
    method __len__ (line 314) | def __len__(self) -> int:
  class StateOKDataset (line 318) | class StateOKDataset(PretrainDataset):
    method __init__ (line 319) | def __init__(
    method set_epoch (line 353) | def set_epoch(self, epoch: int) -> None:
    method __getitem__ (line 407) | def __getitem__(self, index: int) -> Tuple[torch.Tensor, torch.Tensor,...
    method __len__ (line 487) | def __len__(self) -> int:
  class GenStateOKDataset (line 491) | class GenStateOKDataset(PretrainDataset):
    method __init__ (line 492) | def __init__(
    method set_epoch (line 523) | def set_epoch(self, epoch: int) -> None:
    method __getitem__ (line 576) | def __getitem__(
    method __len__ (line 660) | def __len__(self) -> int:
  class QuintetDataset (line 664) | class QuintetDataset(PretrainDataset):
    method __init__ (line 665) | def __init__(
    method set_epoch (line 694) | def set_epoch(self, epoch: int) -> None:
    method __getitem__ (line 747) | def __getitem__(self, index: int) -> Tuple[torch.Tensor, torch.Tensor,...
    method __len__ (line 787) | def __len__(self) -> int:
  function get_epoch_datasets (line 791) | def get_epoch_datasets(

FILE: voltron/models/core/vcond.py
  class VCond (line 28) | class VCond(nn.Module):
    method __init__ (line 29) | def __init__(
    method initialize_weights (line 179) | def initialize_weights(self) -> None:
    method transformer_initializer (line 204) | def transformer_initializer(m: nn.Module) -> None:
    method encode_language (line 214) | def encode_language(self, lang: torch.Tensor, lang_mask: torch.Tensor)...
    method mask (line 221) | def mask(
    method get_representations (line 246) | def get_representations(
    method encode (line 281) | def encode(self, img: torch.Tensor, lang: torch.Tensor, lang_mask: tor...
    method forward_encoder (line 313) | def forward_encoder(
    method forward_decoder (line 349) | def forward_decoder(self, visible_patches: torch.Tensor, restore_idxs:...
    method patchify (line 386) | def patchify(self, imgs: torch.Tensor) -> torch.Tensor:
    method compute_loss (line 395) | def compute_loss(self, imgs: torch.Tensor, reconstructions: torch.Tens...
    method forward (line 410) | def forward(
    method configure_optimizer (line 418) | def configure_optimizer(self) -> Tuple[torch.optim.Optimizer, Callable...

FILE: voltron/models/core/vdual.py
  class VDual (line 28) | class VDual(nn.Module):
    method __init__ (line 29) | def __init__(
    method initialize_weights (line 185) | def initialize_weights(self) -> None:
    method transformer_initializer (line 210) | def transformer_initializer(m: nn.Module) -> None:
    method encode_language (line 220) | def encode_language(self, lang: torch.Tensor, lang_mask: torch.Tensor)...
    method mask (line 227) | def mask(
    method get_representations (line 255) | def get_representations(
    method encode (line 292) | def encode(self, imgs: torch.Tensor, lang: torch.Tensor, lang_mask: to...
    method forward_encoder (line 333) | def forward_encoder(
    method forward_decoder (line 372) | def forward_decoder(self, visible_patches: torch.Tensor, restore_idxs:...
    method patchify (line 424) | def patchify(self, imgs: torch.Tensor) -> torch.Tensor:
    method compute_loss (line 433) | def compute_loss(
    method forward (line 454) | def forward(
    method configure_optimizer (line 477) | def configure_optimizer(self) -> Tuple[torch.optim.Optimizer, Callable...

FILE: voltron/models/core/vgen.py
  class VGen (line 32) | class VGen(nn.Module):
    method __init__ (line 33) | def __init__(
    method initialize_weights (line 210) | def initialize_weights(self) -> None:
    method transformer_initializer (line 235) | def transformer_initializer(m: nn.Module) -> None:
    method embed_language (line 245) | def embed_language(self, lang: torch.Tensor) -> torch.Tensor:
    method encode_language (line 253) | def encode_language(self, lang: torch.Tensor, lang_mask: torch.Tensor)...
    method mask (line 260) | def mask(
    method get_representations (line 288) | def get_representations(
    method encode (line 325) | def encode(self, imgs: torch.Tensor, lang: torch.Tensor, lang_mask: to...
    method score (line 369) | def score(self, imgs: torch.Tensor, langs: torch.Tensor, lang_masks: t...
    method forward_encoder (line 479) | def forward_encoder(
    method forward_decoder (line 524) | def forward_decoder(
    method patchify (line 601) | def patchify(self, imgs: torch.Tensor) -> torch.Tensor:
    method compute_loss (line 610) | def compute_loss(
    method forward (line 677) | def forward(
    method configure_optimizer (line 718) | def configure_optimizer(self) -> Tuple[torch.optim.Optimizer, Callable...

FILE: voltron/models/instantiate.py
  function get_model_optimizer (line 22) | def get_model_optimizer(

FILE: voltron/models/materialize.py
  function available_models (line 66) | def available_models() -> List[str]:
  function load (line 70) | def load(

FILE: voltron/models/reproductions/vmvp.py
  class VMVP (line 21) | class VMVP(nn.Module):
    method __init__ (line 22) | def __init__(
    method initialize_weights (line 113) | def initialize_weights(self) -> None:
    method transformer_initializer (line 131) | def transformer_initializer(m: nn.Module) -> None:
    method mask (line 141) | def mask(
    method get_representations (line 166) | def get_representations(self, img: torch.Tensor, mode: str = "patch") ...
    method encode (line 184) | def encode(self, img: torch.Tensor) -> torch.Tensor:
    method forward_encoder (line 199) | def forward_encoder(
    method forward_decoder (line 221) | def forward_decoder(self, visible_patches: torch.Tensor, restore_idxs:...
    method patchify (line 250) | def patchify(self, imgs: torch.Tensor) -> torch.Tensor:
    method compute_loss (line 259) | def compute_loss(self, imgs: torch.Tensor, reconstructions: torch.Tens...
    method forward (line 274) | def forward(
    method configure_optimizer (line 283) | def configure_optimizer(self) -> Tuple[torch.optim.Optimizer, Callable...

FILE: voltron/models/reproductions/vr3m.py
  class VR3M (line 24) | class VR3M(nn.Module):
    method __init__ (line 25) | def __init__(
    method initialize_weights (line 134) | def initialize_weights(self) -> None:
    method transformer_initializer (line 146) | def transformer_initializer(m: nn.Module) -> None:
    method get_representations (line 156) | def get_representations(self, img: torch.Tensor) -> torch.Tensor:
    method encode_images (line 174) | def encode_images(self, imgs: torch.Tensor) -> torch.Tensor:
    method encode_language (line 187) | def encode_language(self, lang: torch.Tensor, lang_mask: torch.Tensor)...
    method get_reward (line 194) | def get_reward(self, initial: torch.Tensor, later: torch.Tensor, lang:...
    method forward (line 197) | def forward(self, imgs: torch.Tensor, lang: torch.Tensor, lang_mask: t...
    method time_similarity (line 236) | def time_similarity(state_x: torch.Tensor, state_y: torch.Tensor, use_...
    method get_time_contrastive_loss (line 241) | def get_time_contrastive_loss(
    method get_reward_loss (line 281) | def get_reward_loss(
    method configure_optimizer (line 341) | def configure_optimizer(self) -> Tuple[torch.optim.Optimizer, Callable...

FILE: voltron/models/reproductions/vrn3m.py
  class VRN3M (line 24) | class VRN3M(nn.Module):
    method __init__ (line 25) | def __init__(
    method get_representations (line 102) | def get_representations(self, img: torch.Tensor) -> torch.Tensor:
    method encode_images (line 113) | def encode_images(self, imgs: torch.Tensor) -> torch.Tensor:
    method encode_language (line 117) | def encode_language(self, lang: torch.Tensor, lang_mask: torch.Tensor)...
    method get_reward (line 124) | def get_reward(self, initial: torch.Tensor, later: torch.Tensor, lang:...
    method extract_features (line 127) | def extract_features(self, img: torch.Tensor) -> torch.Tensor:
    method forward (line 131) | def forward(self, imgs: torch.Tensor, lang: torch.Tensor, lang_mask: t...
    method time_similarity (line 170) | def time_similarity(state_x: torch.Tensor, state_y: torch.Tensor, use_...
    method get_time_contrastive_loss (line 175) | def get_time_contrastive_loss(
    method get_reward_loss (line 215) | def get_reward_loss(
    method configure_optimizer (line 275) | def configure_optimizer(self) -> Tuple[torch.optim.Optimizer, Callable...

FILE: voltron/models/util/extraction.py
  class MAPAttention (line 22) | class MAPAttention(nn.Module):
    method __init__ (line 23) | def __init__(self, embed_dim: int, n_heads: int) -> None:
    method forward (line 33) | def forward(self, seed: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
  class MAPBlock (line 51) | class MAPBlock(nn.Module):
    method __init__ (line 52) | def __init__(
    method forward (line 88) | def forward(self, x: torch.Tensor) -> torch.Tensor:
  function instantiate_extractor (line 96) | def instantiate_extractor(backbone: nn.Module, n_latents: int = 1) -> Ca...

FILE: voltron/models/util/optimization.py
  function get_lr_update (line 19) | def get_lr_update(

FILE: voltron/models/util/transformer.py
  function get_1D_sine_cosine (line 22) | def get_1D_sine_cosine(dim: int, pos: np.ndarray) -> np.ndarray:
  function get_1D_position_embeddings (line 31) | def get_1D_position_embeddings(embed_dim: int, length: int) -> np.ndarray:
  function get_2D_position_embeddings (line 37) | def get_2D_position_embeddings(embed_dim: int, grid_size: int, cls_token...
  class PatchEmbed (line 57) | class PatchEmbed(nn.Module):
    method __init__ (line 58) | def __init__(
    method forward (line 73) | def forward(self, patches: torch.Tensor) -> torch.Tensor:
  class LayerScale (line 84) | class LayerScale(nn.Module):
    method __init__ (line 85) | def __init__(self, dim: int, init_values: float = 0.1) -> None:  # CaI...
    method forward (line 89) | def forward(self, x: torch.Tensor) -> torch.Tensor:
  class RMSNorm (line 94) | class RMSNorm(nn.Module):
    method __init__ (line 95) | def __init__(self, dim: int, eps: float = 1e-8) -> None:
    method forward (line 100) | def forward(self, x: torch.Tensor) -> torch.Tensor:
  class SwishGLU (line 106) | class SwishGLU(nn.Module):
    method __init__ (line 107) | def __init__(self, in_dim: int, out_dim: int) -> None:
    method forward (line 111) | def forward(self, x: torch.Tensor) -> torch.Tensor:
  class Attention (line 119) | class Attention(nn.Module):
    method __init__ (line 120) | def __init__(self, embed_dim: int, n_heads: int, dropout: float = 0.0)...
    method forward (line 131) | def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None...
  class Block (line 158) | class Block(nn.Module):
    method __init__ (line 159) | def __init__(
    method forward (line 204) | def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None...

FILE: voltron/overwatch/overwatch.py
  class OverwatchRich (line 16) | class OverwatchRich:
  class OverwatchStandard (line 47) | class OverwatchStandard:

FILE: voltron/preprocessing/core.py
  function get_path (line 34) | def get_path(save_dir: Path, v: str, i: int, relpath: bool = False) -> str:
  function do_dry_run (line 41) | def do_dry_run(
  function process_clip (line 110) | def process_clip(
  function serialize_epoch (line 150) | def serialize_epoch(

FILE: voltron/preprocessing/process.py
  function extract_frames (line 32) | def extract_frames(
  function preprocess_language (line 107) | def preprocess_language(
  function unify_batches (line 214) | def unify_batches(

FILE: voltron/preprocessing/transforms.py
  function identity (line 20) | def identity(x: torch.Tensor) -> torch.Tensor:
  function scaled_center_crop (line 24) | def scaled_center_crop(target_resolution: int, frames: List[Image.Image]...
  function get_preprocess_transform (line 40) | def get_preprocess_transform(
  function get_online_transform (line 50) | def get_online_transform(

FILE: voltron/preprocessing/v1/process.py
  function preprocess_videos (line 35) | def preprocess_videos(
  function preprocess_language (line 118) | def preprocess_language(
  function jsonify_language (line 186) | def jsonify_language(train_registry: Path, val_registry: Path) -> None:
  function index (line 224) | def index(train_registry: Path, val_registry: Path, name: str, artifact_...
  function unify_batches (line 244) | def unify_batches(

FILE: voltron/preprocessing/v1/transforms.py
  class ComposeMix (line 15) | class ComposeMix:
    method __init__ (line 16) | def __init__(self, transforms):
    method __call__ (line 19) | def __call__(self, imgs):
  class RandomCropVideo (line 31) | class RandomCropVideo:
    method __init__ (line 32) | def __init__(self, size):
    method __call__ (line 35) | def __call__(self, imgs):
  class Scale (line 44) | class Scale:
    method __init__ (line 45) | def __init__(self, size):
    method __call__ (line 48) | def __call__(self, img):
  function identity (line 52) | def identity(x):
  function get_pre_transform (line 57) | def get_pre_transform(dataset: str, resolution: int, scale_factor: float...
  function get_online_transform (line 79) | def get_online_transform(dataset: str, model_arch: str, normalization: T...

FILE: voltron/preprocessing/v1/utils.py
  function get_path (line 31) | def get_path(save_dir: Path, v: str, i: int) -> str:
  function do_dry_run (line 35) | def do_dry_run(
  function process_video (line 100) | def process_video(
  function precompute_epoch (line 138) | def precompute_epoch(

FILE: voltron/util/checkpointing.py
  class FixedDeck (line 26) | class FixedDeck(deque):
    method __init__ (line 27) | def __init__(self, maxlen: int) -> None:
    method append (line 30) | def append(self, x: Any) -> Any:
  class CheckpointSaver (line 40) | class CheckpointSaver:
    method __init__ (line 41) | def __init__(self, strategy: Tuple[int, int, int], run_dir: str, is_ra...
    method save (line 66) | def save(
  function do_resume (line 114) | def do_resume(resume: bool, run_dir: str) -> Tuple[Optional[Path], int, ...

FILE: voltron/util/metrics.py
  class Logger (line 25) | class Logger(ABC):
    method __init__ (line 26) | def __init__(self, run_id: str, hparams: Dict[str, Any], is_rank_zero:...
    method write_hyperparameters (line 30) | def write_hyperparameters(self) -> None:
    method write (line 34) | def write(self, global_step: int, metrics: Dict[str, Union[int, float]...
    method finalize (line 37) | def finalize(self) -> None:
  class JSONLinesLogger (line 41) | class JSONLinesLogger(Logger):
    method write_hyperparameters (line 42) | def write_hyperparameters(self) -> None:
    method write (line 56) | def write(self, global_step: int, metrics: Dict[str, Union[int, float]...
  class WeightsBiasesLogger (line 65) | class WeightsBiasesLogger(Logger):
    method __init__ (line 66) | def __init__(
    method initialize (line 97) | def initialize(self) -> None:
    method write_hyperparameters (line 115) | def write_hyperparameters(self) -> None:
    method write (line 122) | def write(self, global_step: int, metrics: Dict[str, Union[int, float]...
    method finalize (line 129) | def finalize(self) -> None:
  class Metrics (line 137) | class Metrics:
    method __init__ (line 138) | def __init__(
    method itemize (line 220) | def itemize(self) -> Dict[str, torch.Tensor]:
    method log (line 228) | def log(self, global_step: int, metrics: Dict[str, Union[int, float]])...
    method finalize (line 232) | def finalize(self) -> None:
    method get_status (line 236) | def get_status(self, epoch: int, loss: Optional[torch.Tensor] = None) ...
    method commit (line 244) | def commit(
    method push (line 275) | def push(self, epoch: int) -> str:
    method push_epoch (line 369) | def push_epoch(self, epoch: int, val_loss: torch.Tensor) -> Tuple[str,...

FILE: voltron/util/utilities.py
  function worker_init_function (line 43) | def worker_init_function(worker_id: int) -> None:
  function set_global_seed (line 77) | def set_global_seed(seed: int, get_worker_init_fn: bool = False) -> Opti...
  class ResumeableDistributedSampler (line 93) | class ResumeableDistributedSampler(DistributedSampler):
    method __init__ (line 94) | def __init__(
    method __iter__ (line 111) | def __iter__(self) -> Iterator[T_co]:
    method __len__ (line 120) | def __len__(self) -> int:
    method set_epoch (line 127) | def set_epoch(self, epoch: int) -> None:

FILE: voltron/util/v1/checkpointing.py
  class FixedDeck (line 20) | class FixedDeck(deque):
    method __init__ (line 21) | def __init__(self, maxlen: int) -> None:
    method append (line 24) | def append(self, x: Any) -> Any:
  class XLACheckpointSaver (line 34) | class XLACheckpointSaver:
    method __init__ (line 35) | def __init__(self, strategy: Tuple[int, int, int], run_dir: str) -> None:
    method save (line 61) | def save(

FILE: voltron/util/v1/distributed.py
  class ResumeableDistributedSampler (line 27) | class ResumeableDistributedSampler(DistributedSampler):
    method __init__ (line 28) | def __init__(
    method __iter__ (line 45) | def __iter__(self) -> Iterator[T_co]:
    method __len__ (line 54) | def __len__(self) -> int:
    method set_epoch (line 61) | def set_epoch(self, epoch: int) -> None:
  function xla_available (line 69) | def xla_available() -> bool:
  function get_rank (line 76) | def get_rank() -> int:

FILE: voltron/util/v1/random.py
  function set_global_seed (line 24) | def set_global_seed(seed: int) -> Callable[[int], None]:
  function worker_init_function (line 37) | def worker_init_function(worker_id: int) -> None:

FILE: voltron/util/v1/xla_logger.py
  function log_epoch_end_update (line 17) | def log_epoch_end_update(
  function log_vmvp_train_update (line 67) | def log_vmvp_train_update(
  function log_vr3m_train_update (line 99) | def log_vr3m_train_update(
  function log_vrn3m_train_update (line 145) | def log_vrn3m_train_update(
  function log_vcond_train_update (line 192) | def log_vcond_train_update(
  function log_vdual_train_update (line 224) | def log_vdual_train_update(
  function log_vgen_train_update (line 262) | def log_vgen_train_update(