torch.compile specializes on output name

[msaroufim]

🐛 Describe the bug

x-posting https://discuss.pytorch.org/t/pytorch-compile-requires-fixed-input-naming/181323

import torch
torch._dynamo.config.verbose=True

class MyModule(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.lin = torch.nn.Linear(100, 10)

    def forward(self, x, step):
        return {f'output_{step}': self.lin(x[f'input_{step}'])}

mod = MyModule()
opt_mod = torch.compile(mod)

my_input = {
    'input_0': torch.ones([100]),
    'input_1': torch.ones([100])
}

for step in range(2):
    output = opt_mod(my_input, step)
    print(output.keys())

Error logs

Expected Output

dict_keys(['output_0'])
dict_keys(['output_1'])

Actual Output

dict_keys(['output_0'])
dict_keys(['output_0'])

Minified repro

n/a

Versions

.

cc @ezyang @bdhirsh @anijain2305 @zou3519 @chauhang @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @gchanan @wconstab

[zou3519]

Tentatively marking hi-pri for silent incorrectness

[zou3519]

@anijain2305 we marked this issue as hi-pri and you are in the feature map for Dynamo, so I assigned you as per the PT2 oncall instructions. Could you take a look please?

[voznesenskym]

You need to take a guard on step and we don't

[mlazos]

This has been fixed in the nightly