python3 main_umap.py \
    --dataset imagenet100 \
    --train_data_path ./datasets/imagenet-100/train \
    --val_data_path ./datasets/imagenet-100/val \
    --batch_size 16 \
    --num_workers 10 \
    --pretrained_checkpoint_dir $1


python3 main_umap.py \
    --dataset cifar10 \
    --train_data_path ./datasets/imagenet-100/train \
    --val_data_path ./datasets/imagenet-100/val \
    --batch_size 16 \
    --num_workers 10 \
    --pretrained_checkpoint_dir "/extra_room/junlin/SSL/solo-learn/trained_models/byol/7s3i8i56/"