Fix shape annotations for sigmoid-bce loss.

PiperOrigin-RevId: 553113646
google-research · Aug 2, 2023 · e51f224 · e51f224
1 parent d4b3d8c
commit e51f224
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/kauldron/losses/simple.py b/kauldron/losses/simple.py
@@ -91,5 +91,7 @@ class SigmoidBinaryCrossEntropy(base.Loss):
   labels: Key = "batch.label"
 
   @typechecked
-  def get_values(self, logits: Float["*a"], labels: Float["*a"]) -> Float["*a"]:
+  def get_values(
+      self, logits: Float["*a n"], labels: Int["*a n"]
+  ) -> Float["*a 1"]:
     return optax.sigmoid_binary_cross_entropy(logits, labels)