Loads a checkpoint from a given file into state. Args: state: a mapping contaning model, optimizer and lr scheduler p

Initializes multi-GPU-friendly python command line logger.

:param x: tensor of shape [B (batch-size), N (num-patches), D (feature-dim)] :param masks: list of tensors containing indices of patches in [

Reconstruct an image from patches of a given size and stride. Args: x: Patch tensor of shape (B, patch_size * patch_size, C).

imgs: (N, 3, H, W) x: (N, L, patch_size**2 *3)

We use the slip horizon labels to debounce fast switches in slip events, which is likely noise if there is atleast one slip event in

Adds values as postfix string to progressbar. Args: prog_bar: a progressbar (on global rank zero) or an iterable (every other ran

Rescale image pixels to span range [0, 1]

embed_dim: output dimension for each position pos: a list of positions to be encoded: size (M,) out: (M, D)

embed_dim: output dimension for each position pos: a list of positions to be encoded: size (M,) out: (M, D)

embed_dim: output dimension for each position pos: a list of positions to be encoded: size (M,) out: (M, D)

grid_size: int of the grid height and width return: pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/

this will perform the index select, cat the tensors, and provide the attn_bias from cache

this will perform the index select, cat the tensors, and provide the attn_bias from cache

Calculate gradient in data. Args: data (Tensor): input data with shape (B, C, H, W). stride (int): stride for distance of calcula