Train a Parametrized Quantum Circuit with a loss function using the samples measurement

Hello,
While trying to implement this article Implementable Quantum Classifier for Nonlinear Data i syumble upon a problem.
I cannot find a way to train my circuit using a loss function that is defined with the amplitudes of the final state and not on average measurement over every Qubits.

Here are the relevant parts of my code :

def gaussian_kernel(x,y,sigma) :
return np.exp(-np.abs(x-y)**2/2*sigma)

def loss_func(output,k):
l_outuput = list(output)
dico_res = {}
count = 0
for el in output[0] :
I_reg = el.numpy()[:N_qubit_F]
key = sum([2**i for i in I_reg])
if key not in dico_res.keys():
dico_res[key]=1
else:
dico_res[key]+=1
count += 1
loss = 0
for key in dico_res.keys():
loss += gaussian_kernel(dico_res[key],dico_res[key],sigma)
loss += gaussian_kernel(dico_res[key],0,sigma)
if k in dico_res.keys():
loss -= gaussian_kernel(dico_res[key],0,sigma)
loss += gaussian_kernel(dico_res[key],1,sigma)
loss += 1
return loss/count

model = tf.keras.Sequential([
tf.keras.layers.Input(shape=(), dtype=tf.string),
tfq.layers.PQC(circuit,[cirq.Z(cirq.GridQubit(0,0))])
])

@tf.function
def train_step(circuit_input):
sample_layer = tfq.layers.Sample()
output = sample_layer(circuit_input[0] + circuit, repetitions=100)

``````    loss = loss_func(output,k)
print(loss)