OUTPUT_DIR='checkpoints/3frame_patch8x8_mr0.90_gpu/'
DATA_PATH="cwm/data/video_file_lists/kinetics_400_train_list.txt"
MASTER_ADDRESS=10.102.2.146
NNODES=1
NODE_RANK=0
NPROC_PER_NODE=1

echo "master addr: $MASTER_ADDRESS"
echo "num of nodes: $NNODES"
echo "node rank: $NODE_RANK"
echo "procs per node: $NPROC_PER_NODE"

OMP_NUM_THREADS=1 torchrun \
    --nproc_per_node=$NPROC_PER_NODE --nnodes=$NNODES --node_rank=$NODE_RANK \
    --master_addr=$MASTER_ADDRESS --master_port=19234 \
    cwm/run_pretraining.py \
    --data_path ${DATA_PATH} \
    --model vitb_8x8patch_3frames \
    --mask_type rotated_table \
    --mask_ratio 0.90 \
    --mask_kwargs '{"tube_length": 1}' \
    --context_frames 2 \
    --target_frames 1 \
    --temporal_units 'ms' \
    --sampling_rate 150 \
    --context_target_gap 150 150 \
    --batch_size 1 \
    --accum_iter 1 \
    --opt adamw \
    --opt_betas 0.9 0.95 \
    --warmup_epochs 40 \
    --save_ckpt_freq 50 \
    --epochs 800 \
    --augmentation_type 'multiscale' \
    --augmentation_scales 1.0 0.875 0.75 0.66 \
    --log_dir ${OUTPUT_DIR} \
    --output_dir ${OUTPUT_DIR} \
    --print_freq 1 \
    --num_workers 16