diff --git a/symbolic_nn_tests/experiment2/__init__.py b/symbolic_nn_tests/experiment2/__init__.py
index 82e08f0..fda05b1 100644
--- a/symbolic_nn_tests/experiment2/__init__.py
+++ b/symbolic_nn_tests/experiment2/__init__.py
@@ -70,7 +70,7 @@ def run(tensorboard: bool = True, wandb: bool = True):
         wandb_logger = wandb
 
     test(
-        train_loss=semantic_loss.positive_slope_linear_loss(wandb_logger, version),
+        train_loss=semantic_loss.PositiveSlopeLinearLoss(wandb_logger, version),
         val_loss=unpacking_smooth_l1_loss,
         test_loss=unpacking_smooth_l1_loss,
         version=version,
diff --git a/symbolic_nn_tests/experiment2/model.py b/symbolic_nn_tests/experiment2/model.py
index 7877fb7..215119c 100644
--- a/symbolic_nn_tests/experiment2/model.py
+++ b/symbolic_nn_tests/experiment2/model.py
@@ -82,7 +82,7 @@ def main(
         test_loss=test_loss,
     )
     lmodel.configure_optimizers(optimizer=torch.optim.NAdam, **kwargs)
-    trainer = L.Trainer(max_epochs=10, logger=logger)
+    trainer = L.Trainer(max_epochs=5, logger=logger)
     trainer.fit(model=lmodel, train_dataloaders=train, val_dataloaders=val)
     trainer.test(dataloaders=test)
 
diff --git a/symbolic_nn_tests/experiment2/semantic_loss.py b/symbolic_nn_tests/experiment2/semantic_loss.py
index 5603353..b42f869 100644
--- a/symbolic_nn_tests/experiment2/semantic_loss.py
+++ b/symbolic_nn_tests/experiment2/semantic_loss.py
@@ -18,10 +18,16 @@ import torch
 # proportionality.
 
 
-def positive_slope_linear_loss(wandb_logger=None, version="", device="cuda"):
-    a = nn.Parameter(data=torch.randn(1), requires_grad=True).to(device)
+class PositiveSlopeLinearLoss:
+    def __init__(self, wandb_logger=None, version="", device="cuda", log_freq=50):
+        self.a = nn.Parameter(data=torch.randn(1), requires_grad=True).to(device)
+        self.wandb_logger = wandb_logger
+        self.version = version
+        self.device = device
+        self.log_freq = log_freq
+        self.steps_since_log = 0
 
-    def f(out, y):
+    def __call__(self, out, y):
         x, y_pred = out
         x0, x1 = x
 
@@ -56,14 +62,15 @@ def positive_slope_linear_loss(wandb_logger=None, version="", device="cuda"):
         # We also need to calculate a penalty that incentivizes a positive slope. For this, im using relu
         # to scale the slope as it will penalise negative slopes without just creating a reward hack for
         # maximizing slope.
-        slope_penalty = (nn.functional.relu(a * (-m)) + 1).mean()
+        slope_penalty = (nn.functional.relu(self.a * (-m)) + 1).mean()
 
-        if wandb_logger:
-            wandb_logger.log_metrics({f"{version}-a": a})
+        if self.wandb_logger and (self.steps_since_log >= 50):
+            self.wandb_logger.log_metrics({f"{self.version}-a": self.a})
+            self.steps_since_log = 0
+        else:
+            self.steps_since_log += 1
 
         # Finally, let's get a smooth L1 loss and scale it based on these penalty functions
         return (
             nn.functional.smooth_l1_loss(y_pred, y) * residual_penalty * slope_penalty
         )
-
-    return f