move wandb out of optimizer step (#13)

prescient-design · Sep 23, 2024 · f4918d8 · f4918d8
1 parent e25d8bc
commit f4918d8
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/cortex/optim/generative/_lambo.py b/cortex/optim/generative/_lambo.py
@@ -5,7 +5,6 @@
 import numpy as np
 import pandas as pd
 import torch
-import wandb
 from torch.distributions.kl import kl_divergence
 
 from cortex.attribution import approximate_occlusion
@@ -218,7 +217,7 @@ def step(self) -> None:
             "masked_obj_loss": obj_loss.item(),
             "token_entropy": entropy.item(),
         }
-        wandb.log(metrics)
+        return metrics
 
     def _coordinate_selection(
         self,