TypeError: 'ProbUNet' object is not callable (trying to convert code from Tensorflow 1.10 to 2.2.0)

Question

Using prob_unet(x, y, is_training=True, one_hot_labels=cf.one_hot_labels) in first script, I am trying to call the "class ProbUNet(snt.Module):" in the second scri