Could not find matching concrete function to call from SavedModel

[JHlozek]

I am trying to use this code, however, when trying to run the pipeline on the provided test data, I get the attached error below.
This happens whether I run the command line or the python code.

It seems it could be related to this issue for loading/saving weights as the nmrgnn code uses model.save(): tensorflow/tensorflow#37973