Merge pull request #52 from hijohnnylin/fix_db_runner_assert

FIX: Don't check wandb assert if not using wandb
jbloomAus · Mar 26, 2024 · 5c48811 · 5c48811
2 parents f623ed1 + 1adefda
commit 5c48811
Showing 1 changed file with 12 additions and 11 deletions.
diff --git a/sae_analysis/dashboard_runner.py b/sae_analysis/dashboard_runner.py
@@ -364,19 +364,20 @@ def run(self):
                             step=test_idx,
                         )
 
-        # when done zip the folder
-        shutil.make_archive(self.dashboard_folder, "zip", self.dashboard_folder)
+        if self.use_wandb:
+            # when done zip the folder
+            shutil.make_archive(self.dashboard_folder, "zip", self.dashboard_folder)
 
-        # then upload the zip as an artifact
-        artifact = wandb.Artifact("dashboard", type="zipped_feature_dashboards")
-        artifact.add_file(f"{self.dashboard_folder}.zip")
-        assert run is not None  # keep pyright happy
-        run.log_artifact(artifact)
+            # then upload the zip as an artifact
+            artifact = wandb.Artifact("dashboard", type="zipped_feature_dashboards")
+            artifact.add_file(f"{self.dashboard_folder}.zip")
+            assert run is not None  # keep pyright happy
+            run.log_artifact(artifact)
 
-        # terminate the run
-        run.finish()
+            # terminate the run
+            run.finish()
 
-        # delete the dashboard folder
-        shutil.rmtree(self.dashboard_folder)
+            # delete the dashboard folder
+            shutil.rmtree(self.dashboard_folder)
 
         return