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!