tts_angular training fail assertion in cudagraph_trees

[shunting314]

🐛 Describe the bug

It works if cudagraphs is disabled:

time python benchmarks/dynamo/torchbench.py --backend inductor --amp --performance --only tts_angular --training --disable-cudagraphs

But it fails when cudagraphs is enabled:

time python benchmarks/dynamo/torchbench.py --backend inductor --amp --performance --only tts_angular --training

Full error log: https://gist.github.com/shunting314/e29552a249d50fa08411dc0dfde4c986

Put the last few lines of the log here for convenience:

  File "/home/shunting/ws/pytorch/torch/_inductor/cudagraph_trees.py", line 1757, in run
    out = self._run(new_inputs, function_id)
  File "/home/shunting/ws/pytorch/torch/_inductor/cudagraph_trees.py", line 1810, in _run
    return self.execute_node(child, new_inputs)
  File "/home/shunting/ws/pytorch/torch/_inductor/cudagraph_trees.py", line 1886, in execute_node
    return node.run(new_inputs)
  File "/home/shunting/ws/pytorch/torch/_inductor/cudagraph_trees.py", line 951, in run
    assert data_ptr == new_inputs[idx].data_ptr()
AssertionError

Error logs

No response

Minified repro

No response

Versions

..

cc @mcarilli @ezyang @eellison @msaroufim @bdhirsh @anijain2305 @zou3519

[shunting314]

cc @eellison