Obtain scores from infer_with_gml

Hi, thank you for this project! I was able to set up, train and infer this model on my own (binary classification) dataset.

However, `infer_with_gml` only provides me predicted labels; I would like to obtain a score that I can later use to construct an ROC curve for this model.

Could you guide me on how to modify the code to obtain scores instead of binary predictions?