import torch
import torch.nn as nn

class CustomFunction(torch.autograd.Function, nn.Module):
    def __call__(self, input):
        return self.apply(input)

    @staticmethod
    def forward(ctx, forward_in):
        forward_out = forward_in.clamp(min=0)
        ctx.save_for_backward(forward_in, forward_out)
        return forward_out

    @staticmethod
    def backward(ctx, grad_output):
        forward_in, forward_out = ctx.saved_tensors

        relu_gradients = torch.ones_like(forward_out)
        relu_gradients[forward_in < 0] = 0
        relu_gradients = relu_gradients.mul(grad_output)
        # Some extra functions here
        return relu_gradients


class CustomModel(nn.Module):
    def __init__(self):
        super(CustomModel, self).__init__()
        self.cr = CustomFunction()

    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')


This is the minimal version of the code that causes the following error:
RuntimeError: You attempted to access the anomaly metadata of a custom autograd function but the underlying PyNode has already been deallocated.  The most likely reason this occurred is because you assigned x.grad_fn to a local variable and then let the original variable get deallocated.  Don't do that!  If you really have no way of restructuring your code so this is the case, please file an issue reporting that you are affected by this.

Is there a way to reconstruct the code so that this error isn’t thrown? I don’t understand what the error means.