The implementation is framework agnostic, and will work for PyTorch and TensorFlow models. If we have saved this in
a file named pair_classification.py, we can then import it and register it like this: