Module 'torch.library' has no attribute 'custom_op'

Which Pytorch version are you using? You can check it e.g. via print(torch.__version__) inside your script. I guess you might be using an older release and would need to update.