diff --git a/poetry.lock b/poetry.lock
index bad976e..4fcce23 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -3426,6 +3426,21 @@ files = [
 docs = ["furo (>=2023.7.26)", "proselint (>=0.13)", "sphinx (>=7.1.1)", "sphinx-autodoc-typehints (>=1.24)"]
 test = ["appdirs (==1.4.4)", "covdefaults (>=2.3)", "pytest (>=7.4)", "pytest-cov (>=4.1)", "pytest-mock (>=3.11.1)"]
 
+[[package]]
+name = "plotly"
+version = "5.22.0"
+description = "An open-source, interactive data visualization library for Python"
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "plotly-5.22.0-py3-none-any.whl", hash = "sha256:68fc1901f098daeb233cc3dd44ec9dc31fb3ca4f4e53189344199c43496ed006"},
+    {file = "plotly-5.22.0.tar.gz", hash = "sha256:859fdadbd86b5770ae2466e542b761b247d1c6b49daed765b95bb8c7063e7469"},
+]
+
+[package.dependencies]
+packaging = "*"
+tenacity = ">=6.2.0"
+
 [[package]]
 name = "polars"
 version = "0.20.31"
@@ -4808,6 +4823,21 @@ files = [
     {file = "tbb-2021.12.0-py3-none-win_amd64.whl", hash = "sha256:fc2772d850229f2f3df85f1109c4844c495a2db7433d38200959ee9265b34789"},
 ]
 
+[[package]]
+name = "tenacity"
+version = "8.3.0"
+description = "Retry code until it succeeds"
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "tenacity-8.3.0-py3-none-any.whl", hash = "sha256:3649f6443dbc0d9b01b9d8020a9c4ec7a1ff5f6f3c6c8a036ef371f573fe9185"},
+    {file = "tenacity-8.3.0.tar.gz", hash = "sha256:953d4e6ad24357bceffbc9707bc74349aca9d245f68eb65419cf0c249a1949a2"},
+]
+
+[package.extras]
+doc = ["reno", "sphinx"]
+test = ["pytest", "tornado (>=4.5)", "typeguard"]
+
 [[package]]
 name = "tensorboard"
 version = "2.17.0"
@@ -5485,4 +5515,4 @@ multidict = ">=4.0"
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.11"
-content-hash = "d8cc3168211c9f7eaddf78e15b3077aadb6cda3358dfacfa83da732af83aa899"
+content-hash = "26f68613584a0e1f2a3a8d501e0eaa39fc18197621d8a595bfac6fb4d4455a27"
diff --git a/pyproject.toml b/pyproject.toml
index d77e2ae..e5966e7 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -31,6 +31,7 @@ alive-progress = "^3.1.5"
 hvplot = "^0.10.0"
 pyarrow = "^16.1.0"
 loguru = "^0.7.2"
+plotly = "^5.22.0"
 
 
 [build-system]
diff --git a/symbolic_nn_tests/__main__.py b/symbolic_nn_tests/__main__.py
index 0cdf3c1..42e3b32 100644
--- a/symbolic_nn_tests/__main__.py
+++ b/symbolic_nn_tests/__main__.py
@@ -30,7 +30,7 @@ def main(
     import torch
 
     # Enable tensor cores for compatible GPUs
-    for i in torch.cuda.device_count():
+    for i in range(torch.cuda.device_count()):
         if torch.cuda.get_device_properties(i).major > 6:
             torch.set_float32_matmul_precision("medium")
 
diff --git a/symbolic_nn_tests/experiment2/semantic_loss.py b/symbolic_nn_tests/experiment2/semantic_loss.py
index e8d8ec6..5faaeca 100644
--- a/symbolic_nn_tests/experiment2/semantic_loss.py
+++ b/symbolic_nn_tests/experiment2/semantic_loss.py
@@ -43,27 +43,18 @@ def positive_slope_linear_loss(out, y):
     m, c = linear_fit(diff_electronegativity, y_pred)
 
     # To start with, we want to calculate a penalty based on deviation from a linear relationship
-    # Scaling is being based on 1/sech(w*r) as this increases multiplier as deviation grows.
-    # `w` was selected based on noting that the residual spread before eneg scaling was about 25;
-    # enegs were normalised as x/4, so we want to incentivize a spread of about 25/4~=6, and w=0.2
-    # causes the penalty function to cross 2 at just over 6. Yes, that's a bit arbitrary but we're
-    # just steering the model not applying hard constraints to it shold be fine.
     residual_penalty = (
-        (
-            linear_residuals(diff_electronegativity, y_pred, m, c)
-            / sech(0.2 * diff_electronegativity)
-        )
+        (1 / sech(linear_residuals(diff_electronegativity, y_pred, m, c)))
         .abs()
         .float()
         .mean()
     )
 
-    # We also need to calculate a penalty that incentivizes a positive slope. For this, im using softmax
-    # to scale the slope as it will penalise negative slopes while not just creating a reward hack for
-    # maximizing slope. The softmax function approximates 1 from about 5 onwards, so if we multiply m by
-    # 500, then our penalty should be almost minimised for any slope above 0.01 and maximised below 0.01.
-    # This should suffice for incentivizing the model to favour positive slopes.
-    slope_penalty = (torch.nn.functional.softmax(-m * 500.0) + 1).mean()
+    # We also need to calculate a penalty that incentivizes a positive slope. For this, im using softplus
+    # to scale the slope as it will penalise negative slopes without just creating a reward hack for
+    # maximizing slope.
+    slope_penalty = (nn.functional.softplus(-m) + 1).mean()
 
     # 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
+