diff --git a/bash_files/pretrain/cifar/exe.py b/bash_files/pretrain/cifar/exe.py
index 1bc1697..fcbfdd2 100644
--- a/bash_files/pretrain/cifar/exe.py
+++ b/bash_files/pretrain/cifar/exe.py
@@ -3,6 +3,7 @@
 # poison_method = 'zoo-simclr'
 poison_method = 'clb'
 import time
+out_dir = '/data/yfwang/solo-learn-outputs'
 
 def sweep_poison_rate(args):
     i = 0
@@ -33,10 +34,10 @@ def sweep_pretrain_method(args):
         print(rate)
 
         os.system(f"""
-        for file in /data/yfwang/solo-learn/poison_datasets/{dataset}/{poison_method}/gaussian_noise/{dataset}_{poison_method}_rate_{rate}_*.pt
+        for file in {out_dir}/poison_datasets/{dataset}/{poison_method}/gaussian_noise/{dataset}_{poison_method}_rate_{rate}_*.pt
         do 
         # echo ${{file}}, {method}
-        # CUDA_VISIBLE_DEVICES={gpu} sh {method}.sh {dataset} " --poison_data ${{file}}  --use_poison --checkpoint_dir /data/yfwang/solo-learn/pretrain/{dataset} " &
+        CUDA_VISIBLE_DEVICES={gpu} sh {method}.sh {dataset} " --poison_data ${{file}}  --use_poison --checkpoint_dir {out_dir}/pretrain/{dataset} " &
         done
         """
         )
@@ -59,10 +60,10 @@ def sweep_eval(args):
                     print(rate)
 
                     os.system(f"""
-                    for file in /data/yfwang/solo-learn/poison_datasets/{dataset}/{poison_method}/gaussian_noise/{dataset}_{poison_method}_rate_{rate}_*.pt
+                    for file in {out_dir}/poison_datasets/{dataset}/{poison_method}/gaussian_noise/{dataset}_{poison_method}_rate_{rate}_*.pt
                     do 
                     # echo ${{file}}, {method}
-                    CUDA_VISIBLE_DEVICES={gpu} sh {method}.sh {dataset} " --poison_data ${{file}}  --{apply_method} --checkpoint_dir /data/yfwang/solo-learn/pretrain/{dataset} " &
+                    CUDA_VISIBLE_DEVICES={gpu} sh {method}.sh {dataset} " --poison_data ${{file}}  --{apply_method} --checkpoint_dir {out_dir}/pretrain/{dataset} " &
                     done
                     """
                     )
@@ -84,10 +85,10 @@ def sweep_cifar100(args):
                     print(rate)
 
                     os.system(f"""
-                    for file in /data/yfwang/solo-learn/poison_datasets/{dataset}/{poison_method}/gaussian_noise/{dataset}_{poison_method}_rate_{rate}_*.pt
+                    for file in {out_dir}/poison_datasets/{dataset}/{poison_method}/gaussian_noise/{dataset}_{poison_method}_rate_{rate}_*.pt
                     do 
                     # echo ${{file}}, {method}
-                    CUDA_VISIBLE_DEVICES={gpu} sh {method}.sh {dataset} " --poison_data ${{file}}  --{apply_method} --checkpoint_dir /data/yfwang/solo-learn/pretrain/{dataset} " &
+                    CUDA_VISIBLE_DEVICES={gpu} sh {method}.sh {dataset} " --poison_data ${{file}}  --{apply_method} --checkpoint_dir {out_dir}/pretrain/{dataset} " &
                     done
                     """
                     )
@@ -107,5 +108,5 @@ def sweep_cifar100(args):
     args = parser.parse_args()
     # sweep_pretrain_method(args)
     # sweep_pretrain_method(args)
-    # sweep_eval(args)    
+    # sweep_eval(args)/
     sweep_cifar100(args)
\ No newline at end of file
diff --git a/main_pretrain.py b/main_pretrain.py
index 46394fa..75286e0 100644
--- a/main_pretrain.py
+++ b/main_pretrain.py
@@ -78,6 +78,8 @@ def main():
     if args.num_large_crops != 2:
         assert args.method == "wmse"
 
+    args.checkpoint_dir = os.path.join(args.checkpoint_dir, args.method + poison_suffix)
+
     MethodClass = METHODS[args.method]
     if args.dali:
         assert (
@@ -142,6 +144,7 @@ def main():
             project=args.project,
             entity=args.entity,
             offline=args.offline,
+            save_dir=args.checkpoint_dir,
         )
         wandb_logger.watch(model, log="gradients", log_freq=100)
         wandb_logger.log_hyperparams(args)
@@ -154,7 +157,7 @@ def main():
         # save checkpoint on last epoch only
         ckpt = Checkpointer(
             args,
-            logdir=os.path.join(args.checkpoint_dir, args.method + poison_suffix),
+            logdir=args.checkpoint_dir,
             frequency=args.checkpoint_frequency,
         )
         callbacks.append(ckpt)
@@ -175,7 +178,7 @@ def main():
     ckpt_path = None
     if args.auto_resume and args.resume_from_checkpoint is None:
         auto_resumer = AutoResumer(
-            checkpoint_dir=os.path.join(args.checkpoint_dir, args.method),
+            checkpoint_dir=args.checkpoint_dir,
             max_hours=args.auto_resumer_max_hours,
         )
         resume_from_checkpoint = auto_resumer.find_checkpoint(args)