Obtain scores from infer_with_gml

[emaadmanzoor]

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?