cudagraph error with set_ on inputs

[eellison]

🐛 Describe the bug

We need to track storage liveness for correctness and expect the storages to be on the outputs.

A possible fix here would be to add the inputs which are mutated as outputs, but needs further investigation.

Repro (right now internal only) https://www.internalfb.com/phabricator/paste/view/P1360460962.

Error logs

...

Minified repro

...

Versions

master

cc @mcarilli @ezyang @peterbell10 @msaroufim @bdhirsh @anijain2305 @chauhang

[bdhirsh]

hmm @eellison if we (eventually) care about FSDP2 + cudagraphs, does that mean we're going to have to make sure that the set_() happens outside of the graph? (I think that will only happen if there's a graph break mid-layer, so maybe this isn't too important of an issue..)