take mean for distance-based loss

ChEB-AI · Oct 24, 2023 · 324b9a6 · 324b9a6
1 parent c2947a6
commit 324b9a6
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/chebai/models/electra.py b/chebai/models/electra.py
@@ -543,7 +543,7 @@ def __call__(self, outputs, targets):
         theta = 0.00004
         loss = ( ( torch.sqrt(d) * (d > theta) * t ) + ( (d <= theta) * (1 - d) * (1 - t) ) )
 
-        scalar_loss = torch.max(loss)
+        scalar_loss = torch.mean(loss)
         return scalar_loss
 
 def softabs(x, eps=0.01):