[pt20][eager] Lamb optimizer cannot be used in the compiled function
Labels: triaged, oncall: pt2, module: dynamo

### 🐛 Describe the bug

When I using `torch.compile` to compile a method including the `step` process of the `Lamb` optimizer. It would fail at the second iteration.

### Error logs

```python
[2023-02-08 19:17:53,174] torch._dynamo.variables.torch: [WARNING] Profiler will be ignored
Traceback (most recent call last):
  File "/nvme/mazerun/.conda/envs/torch2.0/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 324, in _compile
    out_code = transform_code_object(code, transform)
  File "/nvme/mazerun/.conda/envs/torch2.0/lib/python3.10/site-packages/torch/_dynamo/bytecode_transformation.py", line 361, in transform_code_object
    transformations(instructions, code_options)
  File "/nvme/mazerun/.conda/envs/torch2.0/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 311, in transform
    tracer.run()
  File "/nvme/mazerun/.conda/envs/torch2.0/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 1683, in run
    super().run()
  File "/nvme/mazerun/.conda/envs/torch2.0/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 569, in run
    and self.step()
  File "/nvme/mazerun/.conda/envs/torch2.0/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 532, in step
    getattr(self, inst.opname)(inst)
  File "/nvme/mazerun/.conda/envs/torch2.0/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 338, in wrapper
    return inner_fn(self, inst)
  File "/nvme/mazerun/.conda/envs/torch2.0/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 144, in impl
    self.push(fn_var.call_function(self, self.popn(nargs), {}))
  File "/nvme/mazerun/.conda/envs/torch2.0/lib/python3.10/site-packages/torch/_dynamo/variables/builtin.py", line 495, in call_function
    result = handler(tx, *args, **kwargs)
  File "/nvme/mazerun/.conda/envs/torch2.0/lib/python3.10/site-packages/torch/_dynamo/variables/builtin.py", line 703, in call_getitem
    return args[0].call_method(tx, "__getitem__", args[1:], kwargs)
  File "/nvme/mazerun/.conda/envs/torch2.0/lib/python3.10/site-packages/torch/_dynamo/variables/dicts.py", line 68, in call_method
    return self.getitem_const(args[0])
  File "/nvme/mazerun/.conda/envs/torch2.0/lib/python3.10/site-packages/torch/_dynamo/variables/dicts.py", line 53, in getitem_const
    return self.items[ConstDictVariable.get_key(arg)].add_options(self, arg)
KeyError: exp_avg_sq

from user code:
   File "/nvme/mazerun/.conda/envs/torch2.0/lib/python3.10/site-packages/timm/optim/lamb.py", line 161, in step
    exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']

Set torch._dynamo.config.verbose=True for more information


You can suppress this exception and fall back to eager by setting:
    torch._dynamo.config.suppress_errors = True


The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/nvme/mazerun/repro.py", line 31, in <module>
    opt(data, optimizer)
  File "/nvme/mazerun/.conda/envs/torch2.0/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 209, in _fn
    return fn(*args, **kwargs)
  File "/nvme/mazerun/repro.py", line 20, in train_step
    loss.backward()
  File "/nvme/mazerun/repro.py", line 21, in <graph break in train_step>
    optimizer.step()
  File "/nvme/mazerun/.conda/envs/torch2.0/lib/python3.10/site-packages/torch/optim/optimizer.py", line 265, in wrapper
    out = func(*args, **kwargs)
  File "/nvme/mazerun/.conda/envs/torch2.0/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/nvme/mazerun/.conda/envs/torch2.0/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 330, in catch_errors
    return callback(frame, cache_size, hooks)
  File "/nvme/mazerun/.conda/envs/torch2.0/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 404, in _convert_frame
    result = inner_convert(frame, cache_size, hooks)
  File "/nvme/mazerun/.conda/envs/torch2.0/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 104, in _fn
    return fn(*args, **kwargs)
  File "/nvme/mazerun/.conda/envs/torch2.0/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 262, in _convert_frame_assert
    return _compile(
  File "/nvme/mazerun/.conda/envs/torch2.0/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 163, in time_wrapper
    r = func(*args, **kwargs)
  File "/nvme/mazerun/.conda/envs/torch2.0/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 394, in _compile
    raise InternalTorchDynamoError() from e
torch._dynamo.exc.InternalTorchDynamoError
```

### Minified repro

```python
from timm.optim import Lamb

import torch
import torch.nn as nn

class Repro(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear1 = nn.Linear(4, 4)
        self.linear2 = nn.Linear(4, 4)

    def forward(self, x):
        x = self.linear1(x)
        x = self.linear2(x)
        return x

    def train_step(self, x, optimizer):
        loss = self(x).mean()
        loss.backward()
        optimizer.step()


if __name__ == "__main__":
    model = Repro().cuda()
    optimizer = Lamb(model.parameters())
    opt = torch.compile(model.train_step, backend='eager')
    data = torch.rand(2, 4).cuda()

    for i in range(2):
        opt(data, optimizer)
```

### Versions

```
timm               0.6.12
torch              2.0.0.dev20230207+cu117
torchvision        0.15.0.dev20230207+cu117
numpy              1.24.2
```

cc @ezyang @soumith @msaroufim @wconstab @ngimel @bdhirsh @voznesenskym @yanboliang @penguinwu @anijain2305 @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @desertfire