fix logits temperature (#868)

Author: sararb

merlin/models/tf/outputs/base.py

@@ -140,7 +140,8 @@ def __call__(self, inputs, *args, **kwargs):
 
         if getattr(self, "logits_scaler", None):
             if isinstance(outputs, tf.Tensor):
-                outputs = Prediction(outputs)
+                targets = kwargs.pop("targets", None)
+                outputs = Prediction(outputs, targets)
             outputs = self.logits_scaler(outputs)
 
         return outputs

tests/unit/tf/models/test_base.py

@@ -860,3 +860,29 @@ def _compute_model_metrics(
     assert not np.isclose(loss1, loss_masked)
     for k in metrics1:
         assert not np.isclose(metrics1[k], metrics_masked[k])
+
+
+def test_categorical_prediction_with_temperature(sequence_testing_data: Dataset):
+    train = sequence_testing_data
+    train.schema = train.schema.select_by_name(["item_id_seq", "user_country"])
+    schema_model = train.schema.select_by_name(["item_id_seq"])
+    inputs = mm.InputBlockV2(
+        schema_model,
+        embeddings=mm.Embeddings(
+            schema_model,
+        ),
+    )
+    model = mm.Model(
+        inputs,
+        mm.MLPBlock([32]),
+        mm.CategoricalOutput(
+            to_call=train.schema.select_by_name(["user_country"]), logits_temperature=0.2
+        ),
+    )
+
+    loader = mm.Loader(
+        train, batch_size=1024, transform=mm.ToTarget(train.schema, "user_country", one_hot=True)
+    )
+
+    model.compile(run_eagerly=False, optimizer="adam")
+    model.fit(loader, batch_size=1024, epochs=1)