import glob
import os
from argparse import ArgumentParser

parser = ArgumentParser()
parser.add_argument('--dir', type=str, required=True, help='Directory to clean up')
args = parser.parse_args()

ckpt_dirs = [dir for dir in glob.glob(os.path.join(args.dir, '*', '*', '*')) if 'checkpoint' in dir]

for dir in ckpt_dirs:
    os.system(f'rm -r {dir}')