search

Lightning-AI / lightning-thunder

Make PyTorch models up to 40% faster! Thunder is a source to source compiler for PyTorch. It enables using different hardware executors at once; across one or thousands of GPUs.
Apache License 2.0
1.07k stars 60 forks source link

pass results through epilogue if present #637

Open t-vi opened 6 days ago

t-vi commented 6 days ago

Here we pass results through the epilogue if we have one. The goal is to have a fix #623 / enhance #632 to put the construction of the dataclass into the epilogue rather than into the compute trace. This way the compute trace could always return a (flat tuple of) plain tensors / numbers.

t-vi commented 6 days ago

I still need to fix the interaction with transformations adding outputs (but that's good, so we can have a principled way for that, too).