Open carmocca opened 3 months ago
You mean something along the lines of Lightning-AI/lit-thunder-LEGACY#1779 ?
Oh yes perfect. I was happy with just not erroring out because otherwise we would need to comment this out in Fabric if we want to compile forward and the loss together
Thinking about this more, it's not so clear that it is a reasonable implementation though, because the JITed things will need backward hook calls generated for them because we don't use the autograd engine for it..
🚀 Feature
Motivation
In Lightning Fabric, we use this once for error checking that the user properly called
backward
. https://github.com/Lightning-AI/pytorch-lightning/blob/096b063d6eeb41567409f4a6b9bac6f5af28ed93/src/lightning/fabric/wrappers.py#L232-L233. cc @awaelchliI don't expect that we run this hook properly on backward, but it would be useful to simply ignore it and not fail on it, maybe showing a warning.
Pitch
Additional context