Updated location of wandb finish()

jbloomAus · Aug 29, 2024 · 921da77 · 921da77
1 parent 3e95ffd
commit 921da77
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/sae_dashboard/neuronpedia/neuronpedia_runner.py b/sae_dashboard/neuronpedia/neuronpedia_runner.py
@@ -8,6 +8,7 @@
 import numpy as np
 import torch
 import wandb
+import wandb.sdk
 from matplotlib import colors
 from sae_lens.sae import SAE
 from sae_lens.toolkit.pretrained_saes import load_sparsity
@@ -485,7 +486,7 @@ def run(self):
                     )
 
         if self.cfg.use_wandb:
-            wandb.finish()
+            wandb.sdk.finish()
 
 
 def main():