I solved this as follows:
#include <torch/script.h>
#include <torch/all.h>

#include <iostream>
#include <memory>


class CustomReluOp : public torch::autograd::Function<CustomReluOp>{
    public:
        static torch::autograd::variable_list forward(torch::autograd::AutogradContext* ctx, torch::autograd::Variable forward_in){
            auto forward_out = torch::clamp(forward_in, 0.0);
            ctx->save_for_backward({forward_in, forward_out});
            return {forward_out};
        }

        static torch::autograd::variable_list backward(torch::autograd::AutogradContext* ctx, torch::autograd::variable_list grad_output){
            auto list_forward = ctx->get_saved_variables();
            auto forward_in = list_forward[0];
            auto forward_out = list_forward[1];

            auto relu_gradients = torch::ones_like(forward_out);
            auto indices = forward_in<0;
            relu_gradients.index({indices}) = 0;
            relu_gradients = torch::mul(relu_gradients, grad_output[0]);

            relu_gradients = torch::nn::functional::relu(relu_gradients);
            relu_gradients = relu_gradients * forward_out;
            return {relu_gradients};
        }
};

torch::Tensor custom_relu_op(const torch::Tensor& input) {
  return CustomReluOp::apply(input)[0];
}

static auto registry = torch::RegisterOperators("my_ops::custom_relu_op", &custom_relu_op);

The code above I saved as customrelu.cpp, I then made the file CMakeLists.txt with the following instructions:
cmake_minimum_required(VERSION 3.1 FATAL_ERROR)
project(custom_relu)

find_package(Torch REQUIRED)

# Define our library target
add_library(custom_relu SHARED customrelu.cpp)
set(CMAKE_CXX_STANDARD 14)
# Link against LibTorch
target_link_libraries(custom_relu "${TORCH_LIBRARIES}")

I then made an empty folder build and cd’d into it and ran the following:
cmake -DCMAKE_PREFIX_PATH="$(python -c 'import torch.utils; print(torch.utils.cmake_prefix_path)')" ..

make -j

The following code now uses a custom operation that you can script:
import torch
import torch.nn as nn

torch.ops.load_library("inference/build/libcustom_relu.so")

class CustomModel(nn.Module):
    def __init__(self):
        super(CustomModel, self).__init__()
        self.cr = torch.ops.my_ops.custom_relu_op

    def forward(self, x):
        x.requires_grad_(True)
        self.cr(x)
        criterion = x**2
        criterion.backward()
        return x.grad

model = CustomModel()
traced_script_module = torch.jit.script(model)
traced_script_module.save("traced_jit_model.pt")
jit_model = torch.jit.load('traced_jit_model.pt')

I added these instructions because the documentation for the c++ frontend is somewhat daunting for someone with little experience with c++ (like me). Credits to this post: TorchScript register backward C++ functions 2