mlpf_evaluate.py

import tensorflow as tf
import numpy as np

#download model_full from https://jpata.web.cern.ch/jpata/2101.08578/v2/model_full.tar.gz
#untar the archive
#note that the custom loss is not included!
model = tf.keras.models.load_model("model_full")
model.summary()

#(batch size, elements in event, features)
inputs = np.random.randn(2, 6400, 12)

outputs = model(inputs)

pred_id = outputs[:, :, :6]
pred_charge = outputs[:, :, 7]
pred_momentum = outputs[:, :, 8:]

print(pred_id)
print(pred_charge)
print(pred_momentum)