[BUG] logits_temperature arg gives error when set in CategoricalOutput class

[rnyak]

Bug description

I get the following error when I set logits_temperature arg in the CategoricalOutput class:

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
File <timed exec>:1

File /usr/local/lib/python3.8/dist-packages/merlin/models/tf/models/base.py:899, in BaseModel.fit(self, x, y, batch_size, epochs, verbose, callbacks, validation_split, validation_data, shuffle, class_weight, sample_weight, initial_epoch, steps_per_epoch, validation_steps, validation_batch_size, validation_freq, max_queue_size, workers, use_multiprocessing, train_metrics_steps, pre, **kwargs)
    896     self._reset_compile_cache()
    897     self.train_pre = pre
--> 899 out = super().fit(**fit_kwargs)
    901 if pre:
    902     del self.train_pre

File /usr/local/lib/python3.8/dist-packages/keras/utils/traceback_utils.py:67, in filter_traceback.<locals>.error_handler(*args, **kwargs)
     65 except Exception as e:  # pylint: disable=broad-except
     66   filtered_tb = _process_traceback_frames(e.__traceback__)
---> 67   raise e.with_traceback(filtered_tb) from None
     68 finally:
     69   del filtered_tb

File /usr/local/lib/python3.8/dist-packages/merlin/models/tf/models/base.py:713, in BaseModel.train_step(self, data)
    710         else:
    711             x = out
--> 713     outputs = self.call_train_test(x, y, sample_weight=sample_weight, training=True)
    714     loss = self.compute_loss(x, outputs.targets, outputs.predictions, outputs.sample_weight)
    716 self._validate_target_and_loss(outputs.targets, loss)

File /usr/local/lib/python3.8/dist-packages/merlin/models/tf/models/base.py:641, in BaseModel.call_train_test(self, x, y, sample_weight, training, testing, **kwargs)
    638     predictions[task.full_name] = task_x
    639     sample_weights[task.full_name] = task_sample_weight
--> 641     self.adjust_predictions_and_targets(predictions, targets)
    643 return Prediction(predictions, targets, sample_weights)

File /usr/local/lib/python3.8/dist-packages/merlin/models/tf/models/base.py:685, in BaseModel.adjust_predictions_and_targets(self, predictions, targets)
    682     predictions[k]._keras_mask = targets[k]._keras_mask
    684 # Ensuring targets and preds have the same dtype
--> 685 targets[k] = tf.cast(targets[k], predictions[k].dtype)
    687 # Ensuring targets are one-hot encoded if they are not
    688 targets[k] = tf.cond(
    689     tf.rank(targets[k]) == tf.rank(predictions[k]) - 1,
    690     lambda: tf.one_hot(
   (...)
    695     lambda: targets[k],
    696 )

ValueError: Attempt to convert a value (None) with an unsupported type (<class 'NoneType'>) to a Tensor.

Steps/Code to reproduce bug

1. First of all, please generate the syntethic data set and process with NVT. you can use this code

2. Once you generate the train set, run the code below to repro the error:

manual_dims = {
    'item_id_list_seq': 64, 
    'item_id_last': 64,
    'f_47_list_seq': 16,
    'f_68_list_seq': 16
}
input_block = mm.InputBlockV2(
        schema_model,
        categorical=mm.Embeddings(schema_model,
                                 dim=manual_dims
                                )
)
mlp_block = mm.MLPBlock(
        [128, 64], 
        no_activation_last_layer=True, 
        dropout=0.2
    )
item_id_name = train.schema.select_by_tag(Tags.ITEM_ID).first.properties['domain']['name']
prediction_task= mm.CategoricalOutput(
        to_call=input_block["categorical"][item_id_name],
        logits_temperature=2,
        target='purchase_id_first'
    )
model_mlp = mm.Model(input_block, mlp_block, prediction_task)
loader = mm.Loader(train, batch_size=1024, transform=mm.ToTarget(train.schema, "purchase_id_first", one_hot=True))
optimizer = tf.keras.optimizers.Adam(
    learning_rate=exp_decay_lr_scheduler,
)

model_mlp.compile(
    optimizer='adam',
    run_eagerly=True,
    loss=tf.keras.losses.CategoricalCrossentropy(
        from_logits=True, 
        label_smoothing=0.2
    ),
    metrics=mm.TopKMetricsAggregator.default_metrics(top_ks=[100])
)
history = model_mlp.fit(loader, epochs=1)

[rnyak]

this was fixed by #868