Add a new compiler backend for jit.trace

Hi all,

I want to add a new compiler backend to pytorch, I saw this post on registering a backend for dynamo.

However, what should I do if I want to do the same for torch.jit.trace? I see that jit.trace does not take any device or backend input argument, but I can find examples like this in the test cases:

traced_net = torch.jit.trace(net, [torch.randn(5, 1, 28, 28, device='cuda')],inp = torch.randn(5, 1, 28, 28, device='cuda')

Is that how we are supposed to get a model trace for a specific backend? Does it mean that I necessarily have to define a new device for that purpose? If so, what should I include in the definition of that device if the only purpose is to use it for jit.trace?

Thank you in advance for helping me!

jit.trace is in deprecated territory, I don’t think its easy to specify and offload to a backend nor will be actively working on it.