fix(sweep): added cuda clean

AdityaNG · May 8, 2024 · 4d3f64f · 4d3f64f
1 parent 921d951
commit 4d3f64f
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/kan_gpt/sweep.py b/kan_gpt/sweep.py
@@ -6,7 +6,6 @@
 
 def wandb_sweep():
     run = wandb.init(resume="allow", anonymous="must")
-    torch.cuda.empty_cache()
 
     class Args:
         model_type = wandb.config.model_type
@@ -21,6 +20,9 @@ class Args:
 
     run_args = Args()
 
+    if "cuda" in run_args.device:
+        torch.cuda.empty_cache()
+
     main(args=run_args, run=run)
 
     if "cuda" in run_args.device: