linting

apax-hub · Jan 9, 2024 · 76bdd87 · 76bdd87
1 parent 9d712f4
commit 76bdd87
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/apax/train/loss.py b/apax/train/loss.py
@@ -52,7 +52,7 @@ def force_angle_exponential_weight(
     return (1.0 - dotp) * jnp.exp(-F_0_norm) / divisor
 
 
-def stress_tril(label, prediction, divisor= 1.0):
+def stress_tril(label, prediction, divisor=1.0):
     idxs = jnp.tril_indices(3)
     label_tril = label[:, idxs[0], idxs[1]]
     prediction_tril = prediction[:, idxs[0], idxs[1]]