[tab] /fast/project/HFMI_ProFOUND/mihir/logs/train/tokenizer/mae/breakout_train_tokenizer_mae_32538.log (log)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=24:00:00
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:1
#SBATCH --mem=100GB
#SBATCH --output=/fast/project/HFMI_ProFOUND/mihir/logs/train/tokenizer/mae/%x_%j.log
#SBATCH --error=/fast/project/HFMI_ProFOUND/mihir/logs/train/tokenizer/mae/%x_%j.log
#SBATCH --job-name=breakout_train_tokenizer_mae

cat $0

source .venv/bin/activate

array_records_dir_train=/fast/project/HFMI_ProFOUND/mihir/data/data_breakout/breakout_episodes_80m/train
array_records_dir_val=/fast/project/HFMI_ProFOUND/mihir/data/data_breakout/breakout_episodes_80m/val

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

CHECKPOINT_DIR=/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/$job_name/$slurm_job_id
mkdir -p $CHECKPOINT_DIR

srun python jasmine/train_tokenizer.py \
    --save_ckpt \
    --image_height=10 \
    --image_width=10 \
    --tokenizer_type=mae \
    --patch_size=4 \
    --ckpt_dir $CHECKPOINT_DIR \
    --name=train_tokenizer_mae_breakout_$slurm_job_id \
    --tags tokenizer mae breakout \
    --entity instant-uv \
    --project jafar \
    --data_dir $array_records_dir_train \
    --val_data_dir $array_records_dir_val \

/var/spool/slurm/job32538/slurm_script: line 15: .venv/bin/activate: No such file or directory
wandb: Currently logged in as: mihir-mahajan2002 (instant-uv) to https://api.wandb.ai. Use `wandb login --relogin` to force relogin
wandb: setting up run zyxrb9ke
wandb: Tracking run with wandb version 0.22.1
wandb: Run data is saved locally in /fast/home/mihir.mahajan/Projects/jasmine_runner_1/wandb/run-20251023_121221-zyxrb9ke
wandb: Run `wandb offline` to turn off syncing.
wandb: Syncing run train_tokenizer_mae_breakout_32538
wandb: ⭐️ View project at https://wandb.ai/instant-uv/jafar
wandb: 🚀 View run at https://wandb.ai/instant-uv/jafar/runs/zyxrb9ke
Running on 1 devices.
Counting all components: ['decoder', 'encoder', 'mask_patch']
Parameter counts:
{'decoder': 16858736, 'encoder': 16858752, 'mask_patch': 48, 'total': 33717536}
WARNING:absl:[process=0][thread=MainThread][operation_id=1] _SignalingThread.join() waiting for signals ([<HandlerAwaitableSignal.STEP_DIRECTORY_CREATION: 'step_directory_creation'>, <HandlerAwaitableSignal.ITEM_DIRECTORY_CREATION: 'item_directory_creation'>]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
WARNING:absl:[process=0][thread=MainThread][operation_id=1] _SignalingThread.join() waiting for signals ([<HandlerAwaitableSignal.STEP_DIRECTORY_CREATION: 'step_directory_creation'>, <HandlerAwaitableSignal.ITEM_DIRECTORY_CREATION: 'item_directory_creation'>]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
WARNING:absl:[process=0][thread=MainThread][operation_id=2] _SignalingThread.join() waiting for signals ([<HandlerAwaitableSignal.STEP_DIRECTORY_CREATION: 'step_directory_creation'>, <HandlerAwaitableSignal.ITEM_DIRECTORY_CREATION: 'item_directory_creation'>]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
WARNING:absl:[process=0][thread=MainThread][operation_id=2] _SignalingThread.join() waiting for signals ([<HandlerAwaitableSignal.STEP_DIRECTORY_CREATION: 'step_directory_creation'>, <HandlerAwaitableSignal.ITEM_DIRECTORY_CREATION: 'item_directory_creation'>]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
WARNING:absl:[process=0][thread=MainThread][operation_id=3] _SignalingThread.join() waiting for signals ([<HandlerAwaitableSignal.STEP_DIRECTORY_CREATION: 'step_directory_creation'>, <HandlerAwaitableSignal.ITEM_DIRECTORY_CREATION: 'item_directory_creation'>]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
WARNING:absl:[process=0][thread=MainThread][operation_id=3] _SignalingThread.join() waiting for signals ([<HandlerAwaitableSignal.STEP_DIRECTORY_CREATION: 'step_directory_creation'>, <HandlerAwaitableSignal.ITEM_DIRECTORY_CREATION: 'item_directory_creation'>]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
WARNING:absl:[process=0][thread=MainThread][operation_id=4] _SignalingThread.join() waiting for signals ([<HandlerAwaitableSignal.STEP_DIRECTORY_CREATION: 'step_directory_creation'>, <HandlerAwaitableSignal.ITEM_DIRECTORY_CREATION: 'item_directory_creation'>]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
WARNING:absl:[process=0][thread=MainThread][operation_id=4] _SignalingThread.join() waiting for signals ([<HandlerAwaitableSignal.STEP_DIRECTORY_CREATION: 'step_directory_creation'>, <HandlerAwaitableSignal.ITEM_DIRECTORY_CREATION: 'item_directory_creation'>]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/004000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/004000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/002000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/002000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/001000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/001000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/003000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/003000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:[process=0][thread=MainThread][operation_id=5] _SignalingThread.join() waiting for signals ([<HandlerAwaitableSignal.STEP_DIRECTORY_CREATION: 'step_directory_creation'>, <HandlerAwaitableSignal.ITEM_DIRECTORY_CREATION: 'item_directory_creation'>]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
WARNING:absl:[process=0][thread=MainThread][operation_id=5] _SignalingThread.join() waiting for signals ([<HandlerAwaitableSignal.STEP_DIRECTORY_CREATION: 'step_directory_creation'>, <HandlerAwaitableSignal.ITEM_DIRECTORY_CREATION: 'item_directory_creation'>]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/004000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/004000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/002000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/002000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/005000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/005000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/003000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/003000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:[process=0][thread=MainThread][operation_id=6] _SignalingThread.join() waiting for signals ([<HandlerAwaitableSignal.STEP_DIRECTORY_CREATION: 'step_directory_creation'>, <HandlerAwaitableSignal.ITEM_DIRECTORY_CREATION: 'item_directory_creation'>]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
WARNING:absl:[process=0][thread=MainThread][operation_id=6] _SignalingThread.join() waiting for signals ([<HandlerAwaitableSignal.STEP_DIRECTORY_CREATION: 'step_directory_creation'>, <HandlerAwaitableSignal.ITEM_DIRECTORY_CREATION: 'item_directory_creation'>]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/004000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/004000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/006000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/006000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/005000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/005000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/003000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/003000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:[process=0][thread=MainThread][operation_id=7] _SignalingThread.join() waiting for signals ([<HandlerAwaitableSignal.STEP_DIRECTORY_CREATION: 'step_directory_creation'>, <HandlerAwaitableSignal.ITEM_DIRECTORY_CREATION: 'item_directory_creation'>]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
WARNING:absl:[process=0][thread=MainThread][operation_id=7] _SignalingThread.join() waiting for signals ([<HandlerAwaitableSignal.STEP_DIRECTORY_CREATION: 'step_directory_creation'>, <HandlerAwaitableSignal.ITEM_DIRECTORY_CREATION: 'item_directory_creation'>]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/004000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/004000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/006000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/006000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/005000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/005000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/007000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/007000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:[process=0][thread=MainThread][operation_id=8] _SignalingThread.join() waiting for signals ([<HandlerAwaitableSignal.STEP_DIRECTORY_CREATION: 'step_directory_creation'>, <HandlerAwaitableSignal.ITEM_DIRECTORY_CREATION: 'item_directory_creation'>]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
WARNING:absl:[process=0][thread=MainThread][operation_id=8] _SignalingThread.join() waiting for signals ([<HandlerAwaitableSignal.STEP_DIRECTORY_CREATION: 'step_directory_creation'>, <HandlerAwaitableSignal.ITEM_DIRECTORY_CREATION: 'item_directory_creation'>]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/006000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/006000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/005000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/005000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/008000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/008000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/007000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/007000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:[process=0][thread=MainThread][operation_id=9] _SignalingThread.join() waiting for signals ([<HandlerAwaitableSignal.STEP_DIRECTORY_CREATION: 'step_directory_creation'>, <HandlerAwaitableSignal.ITEM_DIRECTORY_CREATION: 'item_directory_creation'>]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
WARNING:absl:[process=0][thread=MainThread][operation_id=9] _SignalingThread.join() waiting for signals ([<HandlerAwaitableSignal.STEP_DIRECTORY_CREATION: 'step_directory_creation'>, <HandlerAwaitableSignal.ITEM_DIRECTORY_CREATION: 'item_directory_creation'>]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/006000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/006000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/009000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/009000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/008000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/008000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/007000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/007000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:[process=0][thread=MainThread][operation_id=10] _SignalingThread.join() waiting for signals ([<HandlerAwaitableSignal.STEP_DIRECTORY_CREATION: 'step_directory_creation'>, <HandlerAwaitableSignal.ITEM_DIRECTORY_CREATION: 'item_directory_creation'>]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
WARNING:absl:[process=0][thread=MainThread][operation_id=10] _SignalingThread.join() waiting for signals ([<HandlerAwaitableSignal.STEP_DIRECTORY_CREATION: 'step_directory_creation'>, <HandlerAwaitableSignal.ITEM_DIRECTORY_CREATION: 'item_directory_creation'>]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/010000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/010000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/009000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/009000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/008000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/008000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/007000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/007000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:[process=0][thread=MainThread][operation_id=11] _SignalingThread.join() waiting for signals ([<HandlerAwaitableSignal.STEP_DIRECTORY_CREATION: 'step_directory_creation'>, <HandlerAwaitableSignal.ITEM_DIRECTORY_CREATION: 'item_directory_creation'>]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
WARNING:absl:[process=0][thread=MainThread][operation_id=11] _SignalingThread.join() waiting for signals ([<HandlerAwaitableSignal.STEP_DIRECTORY_CREATION: 'step_directory_creation'>, <HandlerAwaitableSignal.ITEM_DIRECTORY_CREATION: 'item_directory_creation'>]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/010000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/010000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/011000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/011000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/009000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/009000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/008000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/008000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:[process=0][thread=MainThread][operation_id=12] _SignalingThread.join() waiting for signals ([<HandlerAwaitableSignal.STEP_DIRECTORY_CREATION: 'step_directory_creation'>, <HandlerAwaitableSignal.ITEM_DIRECTORY_CREATION: 'item_directory_creation'>]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
WARNING:absl:[process=0][thread=MainThread][operation_id=12] _SignalingThread.join() waiting for signals ([<HandlerAwaitableSignal.STEP_DIRECTORY_CREATION: 'step_directory_creation'>, <HandlerAwaitableSignal.ITEM_DIRECTORY_CREATION: 'item_directory_creation'>]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/010000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/010000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/012000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/012000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/011000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/011000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/009000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/009000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:[process=0][thread=MainThread][operation_id=13] _SignalingThread.join() waiting for signals ([<HandlerAwaitableSignal.STEP_DIRECTORY_CREATION: 'step_directory_creation'>, <HandlerAwaitableSignal.ITEM_DIRECTORY_CREATION: 'item_directory_creation'>]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
WARNING:absl:[process=0][thread=MainThread][operation_id=13] _SignalingThread.join() waiting for signals ([<HandlerAwaitableSignal.STEP_DIRECTORY_CREATION: 'step_directory_creation'>, <HandlerAwaitableSignal.ITEM_DIRECTORY_CREATION: 'item_directory_creation'>]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/010000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/010000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/012000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/012000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/011000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/011000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/013000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/013000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:[process=0][thread=MainThread][operation_id=14] _SignalingThread.join() waiting for signals ([<HandlerAwaitableSignal.STEP_DIRECTORY_CREATION: 'step_directory_creation'>, <HandlerAwaitableSignal.ITEM_DIRECTORY_CREATION: 'item_directory_creation'>]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
WARNING:absl:[process=0][thread=MainThread][operation_id=14] _SignalingThread.join() waiting for signals ([<HandlerAwaitableSignal.STEP_DIRECTORY_CREATION: 'step_directory_creation'>, <HandlerAwaitableSignal.ITEM_DIRECTORY_CREATION: 'item_directory_creation'>]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/014000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/014000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/012000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/012000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/011000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/011000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/013000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/013000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:[process=0][thread=MainThread][operation_id=15] _SignalingThread.join() waiting for signals ([<HandlerAwaitableSignal.STEP_DIRECTORY_CREATION: 'step_directory_creation'>, <HandlerAwaitableSignal.ITEM_DIRECTORY_CREATION: 'item_directory_creation'>]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
WARNING:absl:[process=0][thread=MainThread][operation_id=15] _SignalingThread.join() waiting for signals ([<HandlerAwaitableSignal.STEP_DIRECTORY_CREATION: 'step_directory_creation'>, <HandlerAwaitableSignal.ITEM_DIRECTORY_CREATION: 'item_directory_creation'>]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/014000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/014000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/012000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/012000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/015000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/015000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/013000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/013000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:[process=0][thread=MainThread][operation_id=16] _SignalingThread.join() waiting for signals ([<HandlerAwaitableSignal.STEP_DIRECTORY_CREATION: 'step_directory_creation'>, <HandlerAwaitableSignal.ITEM_DIRECTORY_CREATION: 'item_directory_creation'>]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
WARNING:absl:[process=0][thread=MainThread][operation_id=16] _SignalingThread.join() waiting for signals ([<HandlerAwaitableSignal.STEP_DIRECTORY_CREATION: 'step_directory_creation'>, <HandlerAwaitableSignal.ITEM_DIRECTORY_CREATION: 'item_directory_creation'>]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/014000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/014000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/015000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/015000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/016000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/016000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/013000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/013000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:[process=0][thread=MainThread][operation_id=17] _SignalingThread.join() waiting for signals ([<HandlerAwaitableSignal.STEP_DIRECTORY_CREATION: 'step_directory_creation'>, <HandlerAwaitableSignal.ITEM_DIRECTORY_CREATION: 'item_directory_creation'>]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
WARNING:absl:[process=0][thread=MainThread][operation_id=17] _SignalingThread.join() waiting for signals ([<HandlerAwaitableSignal.STEP_DIRECTORY_CREATION: 'step_directory_creation'>, <HandlerAwaitableSignal.ITEM_DIRECTORY_CREATION: 'item_directory_creation'>]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/014000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/014000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/017000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/017000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/015000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/015000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/016000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/016000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:[process=0][thread=MainThread][operation_id=18] _SignalingThread.join() waiting for signals ([<HandlerAwaitableSignal.STEP_DIRECTORY_CREATION: 'step_directory_creation'>, <HandlerAwaitableSignal.ITEM_DIRECTORY_CREATION: 'item_directory_creation'>]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
WARNING:absl:[process=0][thread=MainThread][operation_id=18] _SignalingThread.join() waiting for signals ([<HandlerAwaitableSignal.STEP_DIRECTORY_CREATION: 'step_directory_creation'>, <HandlerAwaitableSignal.ITEM_DIRECTORY_CREATION: 'item_directory_creation'>]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/018000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/018000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/017000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/017000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/015000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/015000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/016000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/016000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:[process=0][thread=MainThread][operation_id=19] _SignalingThread.join() waiting for signals ([<HandlerAwaitableSignal.STEP_DIRECTORY_CREATION: 'step_directory_creation'>, <HandlerAwaitableSignal.ITEM_DIRECTORY_CREATION: 'item_directory_creation'>]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
WARNING:absl:[process=0][thread=MainThread][operation_id=19] _SignalingThread.join() waiting for signals ([<HandlerAwaitableSignal.STEP_DIRECTORY_CREATION: 'step_directory_creation'>, <HandlerAwaitableSignal.ITEM_DIRECTORY_CREATION: 'item_directory_creation'>]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/018000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/018000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/019000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/019000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/017000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/017000) to end with ".orbax-checkpoint-tmp".
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/016000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/breakout_train_tokenizer_mae/32538/016000) to end with ".orbax-checkpoint-tmp".
Total memory size: 0.8 GB, Output size: 0.4 GB, Temp size: 0.4 GB, Argument size: 0.4 GB, Host temp size: 0.0 GB.
FLOPs: 3.009e+11, Bytes: 1.278e+10 (11.9 GB), Intensity: 23.6 FLOPs/byte
Starting training from step 0...

Memstats: After params initialized:
	Using (GB) 0.4 / 77.61 (0.515398%) on cuda:0
Saved checkpoint at step 1000
Saved checkpoint at step 2000
Saved checkpoint at step 3000
Saved checkpoint at step 4000
Saved checkpoint at step 5000
Saved checkpoint at step 6000
Saved checkpoint at step 7000
Saved checkpoint at step 8000
Saved checkpoint at step 9000
Saved checkpoint at step 10000
Saved checkpoint at step 11000
Saved checkpoint at step 12000
Saved checkpoint at step 13000
Saved checkpoint at step 14000
Saved checkpoint at step 15000
Saved checkpoint at step 16000
Saved checkpoint at step 17000
Saved checkpoint at step 18000
Saved checkpoint at step 19000
Calculating validation metrics...
Traceback (most recent call last):
  File "/fast/home/mihir.mahajan/Projects/jasmine_runner_1/jasmine/train_tokenizer.py", line 611, in <module>
    main(args)
  File "/fast/home/mihir.mahajan/Projects/jasmine_runner_1/jasmine/train_tokenizer.py", line 509, in main
    val_metrics, val_gt_batch, val_recon = calculate_validation_metrics(
                                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/fast/home/mihir.mahajan/Projects/jasmine_runner_1/jasmine/train_tokenizer.py", line 445, in calculate_validation_metrics
    loss, recon, metrics = val_step(tokenizer, batch)
                           ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/fast/home/mihir.mahajan/Projects/jasmine/.venv/lib/python3.12/site-packages/flax/nnx/transforms/compilation.py", line 433, in __call__
    pure_args_out, pure_kwargs_out, pure_out = self.jitted_fn(
                                               ^^^^^^^^^^^^^^^
  File "/fast/home/mihir.mahajan/Projects/jasmine/.venv/lib/python3.12/site-packages/flax/nnx/transforms/compilation.py", line 128, in __call__
    out = self.f(*args, **kwargs)
          ^^^^^^^^^^^^^^^^^^^^^^^
  File "/fast/home/mihir.mahajan/Projects/jasmine_runner_1/jasmine/train_tokenizer.py", line 435, in val_step
    (loss, (recon, metrics)) = tokenizer_loss_fn(tokenizer, inputs, training=False)
                               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/fast/home/mihir.mahajan/Projects/jasmine_runner_1/jasmine/train_tokenizer.py", line 367, in tokenizer_loss_fn
    outputs = model(inputs, training=training)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/fast/home/mihir.mahajan/Projects/jasmine_runner_1/jasmine/models/tokenizer.py", line 206, in __call__
    outputs = self.mask_and_encode(videos_BTHWC, batch["rng"])
                                                 ~~~~~^^^^^^^
KeyError: 'rng'
--------------------
For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.
wandb: 
wandb: 🚀 View run train_tokenizer_mae_breakout_32538 at: https://wandb.ai/instant-uv/jafar/runs/zyxrb9ke
wandb: Find logs at: ../../../../fast/home/mihir.mahajan/Projects/jasmine_runner_1/wandb/run-20251023_121221-zyxrb9ke/logs
W1023 12:25:13.994710   75262 pjrt_client.cc:1469] WatchJobStateAsync failed for task goo.gle/debugonly  job_name: "jax_worker": CANCELLED: CANCELLED
Additional GRPC error information from remote target coordination_service while calling /tensorflow.CoordinationService/WatchJobState:
:UNKNOWN:Error received from peer  {grpc_message:"CANCELLED", grpc_status:1} [type.googleapis.com/tensorflow.DerivedStatus='']
/opt/miniforge3/lib/python3.12/multiprocessing/resource_tracker.py:255: UserWarning: resource_tracker: There appear to be 30 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
srun: error: hai001: task 0: Exited with exit code 1