ValueError in the model initialization

[RoyHirsch]

Hi,
I'm trying to run your code locally and incountered the following error while the model initialization:

ValueError: Output tensors to a Model must be the output of a TensorFlow Layer (thus holding past layer metadata). Found: Tensor("ParticleNet_10/dense_4/Softmax:0", shape=(?, 2), dtype=float32)

I was guessing that maybe one of the operation in the model is not a tf.Layer but i wasn't able to find the solution.
I will appriciate your help.