Tracing models

[jnordberg]

I'm trying to trace the autoregressive and diffusion models using torch but I'm not able to figure out what example inputs to give to make it work. Torch wants only tensors but the methods that need tracing take python arguments and even passes along full dicts, how do you do that with torch.jit.trace_module? (I'm assuming it's possible because the code has support for loading traced models)

[neonbjb]

This was experimental code I forgot to remove. It is possible to add tracing support by fixing some of the model parameters, but it's really tricky. I was able to get it to work for the AR model but there were no performance gains. I haven't figured out the diffusion model yet. Unfortunate since the diffusion model needs it most...

If you manage to figure this out, please submit a PR! I think tracing the diffusion model might get a good performance boost.

[jnordberg]

Ah ok, I was hoping there was some secret incantation to make tracing work with untyped arguments 😄 I'll see if I can get something working by moving all non tensor params into a separate config call, thanks!