refactor: change power loss calculation approach (?)

vsheg · Jul 22, 2024 · f4a8068 · f4a8068
1 parent f84cc13
commit f4a8068
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/moll/measures/_loss.py b/moll/measures/_loss.py
@@ -23,7 +23,7 @@ def power(diff: ArrayLike, p: float = 2.0) -> Array:
         Array(1000., dtype=...)
     """
     diff = jnp.asarray(diff)
-    return jnp.where(diff > 0, jnp.power(diff, p), jnp.inf)
+    return jnp.power(jnp.abs(diff), p) * jnp.sign(diff)  # TODO: is this a good idea?
 
 
 @public