🐛 [Bug] Llama2_flashinfer_rmsnorm example is broken

Logs:

  File "/usr/local/lib/python3.12/dist-packages/torch/_ops.py", line 524, in __call__
    return wrapper()
           ^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/_ops.py", line 515, in wrapper
    return torch.overrides.handle_torch_function(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/overrides.py", line 1725, in handle_torch_function
    result = mode.__torch_function__(public_api, types, args, kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/_export/non_strict_utils.py", line 1035, in __torch_function__
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/_functorch/autograd_function.py", line 49, in __call__
    return super().__call__(autograd_function, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/_ops.py", line 524, in __call__
    return wrapper()
           ^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/_ops.py", line 520, in wrapper
    return self.dispatch(
           ^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/_ops.py", line 481, in dispatch
    assert type(curr_mode) in self.python_key_table, (
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AssertionError: Current active mode <torch.fx.experimental.proxy_tensor.ProxyTorchDispatchMode object at 0x7f2c389c3bc0> not registered