From f4918d835dddac308d373d62b08b3c9ccdea59d0 Mon Sep 17 00:00:00 2001
From: Samuel Stanton <ss13641@nyu.edu>
Date: Mon, 23 Sep 2024 14:07:25 -0400
Subject: [PATCH] move wandb out of optimizer step (#13)

---
 cortex/optim/generative/_lambo.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/cortex/optim/generative/_lambo.py b/cortex/optim/generative/_lambo.py
index 1abcf8f..14d1833 100644
--- 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,