How to use this layer with KerasTensor

[cahuja1992]

I am trying to plugin this layer in Keras model. But it is failing with the exception

*** TypeError: Could not build a TypeSpec for KerasTensor(type_spec=TensorSpec(shape=(None, 12), dtype=tf.float32, name=None), name='dense_1/BiasAdd:0', description="created by layer 'dense_1'") of unsupported type <class 'keras.engine.keras_tensor.KerasTensor'>.

Code to reproduce this issue

import tensorflow as tf
from fast_soft_sort.tf_ops import soft_sort

all_inputs = tf.keras.layers.Input(shape=(12))
x = tf.keras.layers.Dense(12)(all_inputs)
x = soft_sort(x)

model = tf.keras.Model(inputs=all_inputs, outputs=[x])

[yogurt1233]

have you solved this?