The Inference Program might need a hyperparameter or change the base class entirely to GumblePMD. Also, please check whether changing the base class to GumblePMD affects this https://github.com/HLR/DomiKnowS/tree/main/test_regr/examples/PMDExistL/relation_learning_tests.