Repository: pytorch
Issue: Handling custom context managers
Labels: triaged,oncall: pt2,module: dynamo

### 🐛 Describe the bug

The coverage of custom context managers doesn't seem to be very extensive at the moment. With tensordict, we often end up on 

## Context

1. **TD as a context manager**
TensorDict relies on context managers for many "temporary" operations. If, for a period of time, it is convenient to represet a tensordict as a flat structure for example, you could do
```python
td = TensorDict({"a": 0, "b": {"c": 1"}})
with td.flatten_keys() as td_flat:
    c = td_flat["b.c"]
assert c == td["b", "c"]
```

The ops that work in this way are `permute`, `flatten_keys`, `transpose` and a bunch of others. The most used is `to_module` which allows us to make functional calls:
```python
td = TensorDict.from_module(module)
# eg, detach the parameters and cast them in the module
with td.detach().to_module(module):
    output = module(input)
assert not output.requires_grad
```

2. **TensorDictParams**
Similar to [`ParameterList`](https://pytorch.org/docs/stable/generated/torch.nn.ParameterList.html), we use `TensorDictParams` to represent group tensors that are parameters together in a single data structure. The main purpose is that setting such object as an attribute of an nn.Module exposes the content of the tensordict to the nn.Module methods, such as `buffers()`, `parameters()` or anything accessed through `_apply`.

3. **Non-tensor data**
It is useful for tensordict to carry non-tensor data. For one, it is a requirement whenever you want to use tensordict instead of a plain state dict (as the `state_dict()` signature is `Dict[str, Any]`).
Secondly, using tensordict for data preproc requires you to carry non-tensor data along the way (eg, you could store a prompt of a list of prompts along with the tokenized version).

4. **Locking tensordicts**
There is an option to lock/unlock a tensordict. a "locked" tensordict is read-only: you can modify tensors in-place but you will not be able to add or remove an entry. This helps us cache the results of some operations in eager mode if they're cheap to store, and it's also a security when a content is shared between nodes or workers (such that users don't expect that writing a new entry will automatically synchronize the tensordict across workers for example). 

## Where it breaks

Combining 1, 2, 3 and 4 can lead to graph breaks with lengthly error stacks. The best way to reproduce these is to install the tensordict nightly built. I'm reporting a bunch of these in this issue, but I'm also happy to split it in several issues if that's easier to track.

Let's start with what works
```python
from tensordict import from_module, TensorDictParams, TensorDict
import torch.nn

module = torch.nn.Module()
module.params = torch.nn.Parameter(torch.randn(3))
params2 = from_module(module, as_module=True) # as_module produces a TensorDictParams

@torch.compile(fullgraph=True)
def func(z, params2):
    with params2.to_module(module):
        out = z + module.params
    return out

print(func(torch.zeros(()), params2))
```

We can work with a simple tensordict and it works too under #141376
```python

# Isolate the inner tensordict
params2 = params2._param_td
print(func(torch.zeros(()), params2))
```
(here after I assume that tests are executed on that branch)

Now where it breaks: If a module contains a TensorDictParams instance, things start to fall apart
```python
from tensordict import from_module, TensorDictParams, TensorDict
import torch.nn

module = torch.nn.Module()
module.params = TensorDictParams(
    TensorDict(a=0.0)
)
params2 = from_module(module, as_module=True)

@torch.compile(fullgraph=True)
def func(z, params2):
    with params2.to_module(module):
        out = z + module.params["a"]
    return out

print(func(torch.zeros(()), params2))
```

<details>
  <summary>To reproduce</summary>

```
$ pip3 install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cpu -U
$ pip install tensordict-nightly
$ python -c """<code above>"""
```
</details>

Error: see below

Working with a plain tensordict works fine
<details>
  <summary>This works ok</summary>

```python

from tensordict import from_module, TensorDictParams, TensorDict
import torch.nn

module = torch.nn.Module()
module.params = TensorDict(a=0.0)
params2 = from_module(module, as_module=True)

@torch.compile(fullgraph=True)
def func(z, params2):
    with params2.to_module(module):
        out = z + module.params["a"]
    return out

print(func(torch.zeros(()), params2))
```

</details>

In TorchRL, we want to use TensorDictParams to register distribution parameters as buffers such that their device is updated when the policy is sent to a given device (the distribution is built on-the-fly during the forward call of the policy).
Here is how it's done:
https://github.com/pytorch/rl/blob/b7840d7dcf8f67179be2641fb67ee6d2a213418b/sota-implementations/cql/utils.py#L223-L232
in [this PR](https://github.com/pytorch/rl/pull/2553)

Here for some reason, things work ok until we have a `NonTensorData` registered in the distribution kwargs, where the following graph break occurs:

<details>
  <summary>To reproduce</summary>

```
$ git clone https://github.com/pytorch/rl
$ cd rl
$ pip3 install gymnasium wandb tqdm
$ pip3 install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cpu -U
$ pip install tensordict-nightly
$ ghstack checkout https://github.com/pytorch/rl/pull/2553
$ python setup.py develop
$ TORCH_LOGS="+graph_breaks" TORCHDYNAMO_VERBOSE=1 python sota-implementations/cql/cql_online.py compile.compile=1
```

</details>

The funny thing is that this example is actually a more convoluted version of the previous one but things work ok until we have non-tensor data.



### Error logs

<details>
  <summary>Error 1</summary>

```

Traceback (most recent call last):
  File "torch/_dynamo/symbolic_convert.py", line 964, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 661, in wrapper
    return inner_fn(self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 1738, in CALL_FUNCTION_EX
    self.call_function(fn, argsvars.items, kwargsvars)
  File "torch/_dynamo/symbolic_convert.py", line 899, in call_function
    self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
  File "torch/_dynamo/variables/functions.py", line 389, in call_function
    return super().call_function(tx, args, kwargs)
  File "torch/_dynamo/variables/functions.py", line 328, in call_function
    return super().call_function(tx, args, kwargs)
  File "torch/_dynamo/variables/functions.py", line 129, in call_function
    return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 905, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3070, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3196, in inline_call_
    tracer.run()
  File "torch/_dynamo/symbolic_convert.py", line 1054, in run
    while self.step():
  File "torch/_dynamo/symbolic_convert.py", line 969, in step
    self.exception_handler(e)
  File "torch/_dynamo/symbolic_convert.py", line 1554, in exception_handler
    raise raised_exception
  File "torch/_dynamo/symbolic_convert.py", line 964, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 661, in wrapper
    return inner_fn(self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 1738, in CALL_FUNCTION_EX
    self.call_function(fn, argsvars.items, kwargsvars)
  File "torch/_dynamo/symbolic_convert.py", line 899, in call_function
    self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
  File "torch/_dynamo/variables/functions.py", line 328, in call_function
    return super().call_function(tx, args, kwargs)
  File "torch/_dynamo/variables/functions.py", line 129, in call_function
    return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 905, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3070, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3196, in inline_call_
    tracer.run()
  File "torch/_dynamo/symbolic_convert.py", line 1054, in run
    while self.step():
  File "torch/_dynamo/symbolic_convert.py", line 969, in step
    self.exception_handler(e)
  File "torch/_dynamo/symbolic_convert.py", line 1554, in exception_handler
    raise raised_exception
  File "torch/_dynamo/symbolic_convert.py", line 964, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 661, in wrapper
    return inner_fn(self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 1750, in CALL_FUNCTION_KW
    self.call_function(fn, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 899, in call_function
    self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
  File "torch/_dynamo/variables/functions.py", line 389, in call_function
    return super().call_function(tx, args, kwargs)
  File "torch/_dynamo/variables/functions.py", line 328, in call_function
    return super().call_function(tx, args, kwargs)
  File "torch/_dynamo/variables/functions.py", line 129, in call_function
    return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 905, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3070, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3196, in inline_call_
    tracer.run()
  File "torch/_dynamo/symbolic_convert.py", line 1054, in run
    while self.step():
  File "torch/_dynamo/symbolic_convert.py", line 969, in step
    self.exception_handler(e)
  File "torch/_dynamo/symbolic_convert.py", line 1554, in exception_handler
    raise raised_exception
  File "torch/_dynamo/symbolic_convert.py", line 964, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 661, in wrapper
    return inner_fn(self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 1660, in CALL_FUNCTION
    self.call_function(fn, args, {})
  File "torch/_dynamo/symbolic_convert.py", line 899, in call_function
    self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
  File "torch/_dynamo/variables/functions.py", line 129, in call_function
    return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 905, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3070, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3196, in inline_call_
    tracer.run()
  File "torch/_dynamo/symbolic_convert.py", line 1054, in run
    while self.step():
  File "torch/_dynamo/symbolic_convert.py", line 969, in step
    self.exception_handler(e)
  File "torch/_dynamo/symbolic_convert.py", line 1554, in exception_handler
    raise raised_exception
  File "torch/_dynamo/symbolic_convert.py", line 964, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 661, in wrapper
    return inner_fn(self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 1750, in CALL_FUNCTION_KW
    self.call_function(fn, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 899, in call_function
    self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
  File "torch/_dynamo/variables/functions.py", line 389, in call_function
    return super().call_function(tx, args, kwargs)
  File "torch/_dynamo/variables/functions.py", line 328, in call_function
    return super().call_function(tx, args, kwargs)
  File "torch/_dynamo/variables/functions.py", line 129, in call_function
    return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 905, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3070, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3196, in inline_call_
    tracer.run()
  File "torch/_dynamo/symbolic_convert.py", line 1054, in run
    while self.step():
  File "torch/_dynamo/symbolic_convert.py", line 969, in step
    self.exception_handler(e)
  File "torch/_dynamo/symbolic_convert.py", line 1554, in exception_handler
    raise raised_exception
  File "torch/_dynamo/symbolic_convert.py", line 964, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 1446, in RAISE_VARARGS
    self._raise_exception_variable(inst)
  File "torch/_dynamo/symbolic_convert.py", line 1439, in _raise_exception_variable
    raise exc.ObservedException(f"raised exception {val}")
torch._dynamo.exc.ObservedException: raised exception ExceptionVariable()

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "torch/_dynamo/variables/ctx_manager.py", line 174, in exit
    ).call_function(
  File "torch/_dynamo/variables/functions.py", line 389, in call_function
    return super().call_function(tx, args, kwargs)
  File "torch/_dynamo/variables/functions.py", line 328, in call_function
    return super().call_function(tx, args, kwargs)
  File "torch/_dynamo/variables/functions.py", line 129, in call_function
    return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 905, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3070, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3196, in inline_call_
    tracer.run()
  File "torch/_dynamo/symbolic_convert.py", line 1054, in run
    while self.step():
  File "torch/_dynamo/symbolic_convert.py", line 964, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 661, in wrapper
    return inner_fn(self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 1660, in CALL_FUNCTION
    self.call_function(fn, args, {})
  File "torch/_dynamo/symbolic_convert.py", line 899, in call_function
    self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
  File "torch/_dynamo/variables/functions.py", line 328, in call_function
    return super().call_function(tx, args, kwargs)
  File "torch/_dynamo/variables/functions.py", line 129, in call_function
    return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 905, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3070, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3196, in inline_call_
    tracer.run()
  File "torch/_dynamo/symbolic_convert.py", line 1054, in run
    while self.step():
  File "torch/_dynamo/symbolic_convert.py", line 969, in step
    self.exception_handler(e)
  File "torch/_dynamo/symbolic_convert.py", line 1514, in exception_handler
    unimplemented(
  File "torch/_dynamo/exc.py", line 317, in unimplemented
    raise Unsupported(msg, case_name=case_name)
torch._dynamo.exc.Unsupported: exception is raised when top of the block stack is not exception handler (e.g. try .. with .. except). Current TOS is Instruction(opcode=143, opname='SETUP_WITH', arg=21, argval=66, offset=22, starts_line=333, is_jump_target=True, positions=None, target=Instruction(opcode=49, opname='WITH_EXCEPT_START', arg=None, argval=None, offset=66, starts_line=333, is_jump_target=True, positions=None, target=None, exn_tab_entry=None), exn_tab_entry=None)

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "torch/_dynamo/symbolic_convert.py", line 661, in wrapper
    return inner_fn(self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 1660, in CALL_FUNCTION
    self.call_function(fn, args, {})
  File "torch/_dynamo/symbolic_convert.py", line 899, in call_function
    self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
  File "torch/_dynamo/variables/ctx_manager.py", line 1222, in call_function
    return self.ctx.exit(tx, *args)
  File "torch/_dynamo/variables/ctx_manager.py", line 184, in exit
    unimplemented(
  File "torch/_dynamo/exc.py", line 316, in unimplemented
    raise Unsupported(msg, case_name=case_name) from from_exc
torch._dynamo.exc.Unsupported: Unsupported context manager TensorDict(
...)'s __exit__ function

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/Users/vmoens/Library/Application Support/JetBrains/PyCharm2023.3/scratches/scratch_11.py", line 43, in <module>
    print(func(torch.zeros(()), params2))
  File "torch/_dynamo/eval_frame.py", line 573, in _fn
    return fn(*args, **kwargs)
  File "torch/_dynamo/convert_frame.py", line 1379, in __call__
    return self._torchdynamo_orig_callable(
  File "torch/_dynamo/convert_frame.py", line 547, in __call__
    return _compile(
  File "torch/_dynamo/convert_frame.py", line 986, in _compile
    guarded_code = compile_inner(code, one_graph, hooks, transform)
  File "torch/_dynamo/convert_frame.py", line 715, in compile_inner
    return _compile_inner(code, one_graph, hooks, transform)
  File "torch/_utils_internal.py", line 95, in wrapper_function
    return function(*args, **kwargs)
  File "torch/_dynamo/convert_frame.py", line 750, in _compile_inner
    out_code = transform_code_object(code, transform)
  File "torch/_dynamo/bytecode_transformation.py", line 1349, in transform_code_object
    transformations(instructions, code_options)
  File "torch/_dynamo/convert_frame.py", line 231, in _fn
    return fn(*args, **kwargs)
  File "torch/_dynamo/convert_frame.py", line 662, in transform
    tracer.run()
  File "torch/_dynamo/symbolic_convert.py", line 2866, in run
    super().run()
  File "torch/_dynamo/symbolic_convert.py", line 1054, in run
    while self.step():
  File "torch/_dynamo/symbolic_convert.py", line 964, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 667, in wrapper
    unimplemented("Graph break under GenericContextWrappingVariable")
  File "torch/_dynamo/exc.py", line 317, in unimplemented
    raise Unsupported(msg, case_name=case_name)
torch._dynamo.exc.Unsupported: Graph break under GenericContextWrappingVariable

from user code:
   File "/Users/vmoens/Library/Application Support/JetBrains/PyCharm2023.3/scratches/scratch_11.py", line 36, in func
    with params2.to_module(module):

Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information


You can suppress this exception and fall back to eager by setting:
    import torch._dynamo
    torch._dynamo.config.suppress_errors = True

```

</details>

TorchRL example:

<details>
  <summary>Error 2</summary>

```
  File "torch/_dynamo/variables/ctx_manager.py", line 174, in exit
    ).call_function(
  File "torch/_dynamo/variables/functions.py", line 389, in call_function
    return super().call_function(tx, args, kwargs)
  File "torch/_dynamo/variables/functions.py", line 328, in call_function
    return super().call_function(tx, args, kwargs)
  File "torch/_dynamo/variables/functions.py", line 129, in call_function
    return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 905, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3069, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3195, in inline_call_
    tracer.run()
  File "torch/_dynamo/symbolic_convert.py", line 1054, in run
    while self.step():
  File "torch/_dynamo/symbolic_convert.py", line 964, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 661, in wrapper
    return inner_fn(self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 1660, in CALL_FUNCTION
    self.call_function(fn, args, {})
  File "torch/_dynamo/symbolic_convert.py", line 899, in call_function
    self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
  File "torch/_dynamo/variables/functions.py", line 328, in call_function
    return super().call_function(tx, args, kwargs)
  File "torch/_dynamo/variables/functions.py", line 129, in call_function
    return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 905, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3069, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3195, in inline_call_
    tracer.run()
  File "torch/_dynamo/symbolic_convert.py", line 1054, in run
    while self.step():
  File "torch/_dynamo/symbolic_convert.py", line 964, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 661, in wrapper
    return inner_fn(self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 1738, in CALL_FUNCTION_EX
    self.call_function(fn, argsvars.items, kwargsvars)
  File "torch/_dynamo/symbolic_convert.py", line 899, in call_function
    self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
  File "torch/_dynamo/variables/functions.py", line 389, in call_function
    return super().call_function(tx, args, kwargs)
  File "torch/_dynamo/variables/functions.py", line 328, in call_function
    return super().call_function(tx, args, kwargs)
  File "torch/_dynamo/variables/functions.py", line 129, in call_function
    return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 905, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3069, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3195, in inline_call_
    tracer.run()
  File "torch/_dynamo/symbolic_convert.py", line 1054, in run
    while self.step():
  File "torch/_dynamo/symbolic_convert.py", line 964, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 661, in wrapper
    return inner_fn(self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 1738, in CALL_FUNCTION_EX
    self.call_function(fn, argsvars.items, kwargsvars)
  File "torch/_dynamo/symbolic_convert.py", line 899, in call_function
    self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
  File "torch/_dynamo/variables/functions.py", line 328, in call_function
    return super().call_function(tx, args, kwargs)
  File "torch/_dynamo/variables/functions.py", line 129, in call_function
    return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 905, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3069, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3195, in inline_call_
    tracer.run()
  File "torch/_dynamo/symbolic_convert.py", line 1054, in run
    while self.step():
  File "torch/_dynamo/symbolic_convert.py", line 964, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 661, in wrapper
    return inner_fn(self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 1750, in CALL_FUNCTION_KW
    self.call_function(fn, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 899, in call_function
    self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
  File "torch/_dynamo/variables/functions.py", line 389, in call_function
    return super().call_function(tx, args, kwargs)
  File "torch/_dynamo/variables/functions.py", line 328, in call_function
    return super().call_function(tx, args, kwargs)
  File "torch/_dynamo/variables/functions.py", line 129, in call_function
    return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 905, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3069, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3195, in inline_call_
    tracer.run()
  File "torch/_dynamo/symbolic_convert.py", line 1054, in run
    while self.step():
  File "torch/_dynamo/symbolic_convert.py", line 964, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 661, in wrapper
    return inner_fn(self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 1750, in CALL_FUNCTION_KW
    self.call_function(fn, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 899, in call_function
    self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
  File "torch/_dynamo/variables/functions.py", line 389, in call_function
    return super().call_function(tx, args, kwargs)
  File "torch/_dynamo/variables/functions.py", line 328, in call_function
    return super().call_function(tx, args, kwargs)
  File "torch/_dynamo/variables/functions.py", line 129, in call_function
    return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 905, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3069, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3195, in inline_call_
    tracer.run()
  File "torch/_dynamo/symbolic_convert.py", line 1054, in run
    while self.step():
  File "torch/_dynamo/symbolic_convert.py", line 964, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 661, in wrapper
    return inner_fn(self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 1750, in CALL_FUNCTION_KW
    self.call_function(fn, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 899, in call_function
    self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
  File "torch/_dynamo/variables/functions.py", line 389, in call_function
    return super().call_function(tx, args, kwargs)
  File "torch/_dynamo/variables/functions.py", line 328, in call_function
    return super().call_function(tx, args, kwargs)
  File "torch/_dynamo/variables/functions.py", line 129, in call_function
    return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 905, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3069, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3195, in inline_call_
    tracer.run()
  File "torch/_dynamo/symbolic_convert.py", line 1054, in run
    while self.step():
  File "torch/_dynamo/symbolic_convert.py", line 964, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 661, in wrapper
    return inner_fn(self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 1750, in CALL_FUNCTION_KW
    self.call_function(fn, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 899, in call_function
    self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
  File "torch/_dynamo/variables/functions.py", line 389, in call_function
    return super().call_function(tx, args, kwargs)
  File "torch/_dynamo/variables/functions.py", line 328, in call_function
    return super().call_function(tx, args, kwargs)
  File "torch/_dynamo/variables/functions.py", line 129, in call_function
    return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 905, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3069, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3195, in inline_call_
    tracer.run()
  File "torch/_dynamo/symbolic_convert.py", line 1054, in run
    while self.step():
  File "torch/_dynamo/symbolic_convert.py", line 964, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 661, in wrapper
    return inner_fn(self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 1738, in CALL_FUNCTION_EX
    self.call_function(fn, argsvars.items, kwargsvars)
  File "torch/_dynamo/symbolic_convert.py", line 899, in call_function
    self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
  File "torch/_dynamo/variables/functions.py", line 389, in call_function
    return super().call_function(tx, args, kwargs)
  File "torch/_dynamo/variables/functions.py", line 328, in call_function
    return super().call_function(tx, args, kwargs)
  File "torch/_dynamo/variables/functions.py", line 129, in call_function
    return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 905, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3069, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3195, in inline_call_
    tracer.run()
  File "torch/_dynamo/symbolic_convert.py", line 1054, in run
    while self.step():
  File "torch/_dynamo/symbolic_convert.py", line 964, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 661, in wrapper
    return inner_fn(self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 1660, in CALL_FUNCTION
    self.call_function(fn, args, {})
  File "torch/_dynamo/symbolic_convert.py", line 899, in call_function
    self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
  File "torch/_dynamo/variables/functions.py", line 389, in call_function
    return super().call_function(tx, args, kwargs)
  File "torch/_dynamo/variables/functions.py", line 328, in call_function
    return super().call_function(tx, args, kwargs)
  File "torch/_dynamo/variables/functions.py", line 129, in call_function
    return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 905, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3069, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3195, in inline_call_
    tracer.run()
  File "torch/_dynamo/symbolic_convert.py", line 1054, in run
    while self.step():
  File "torch/_dynamo/symbolic_convert.py", line 964, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 661, in wrapper
    return inner_fn(self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 1750, in CALL_FUNCTION_KW
    self.call_function(fn, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 899, in call_function
    self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
  File "torch/_dynamo/variables/functions.py", line 389, in call_function
    return super().call_function(tx, args, kwargs)
  File "torch/_dynamo/variables/functions.py", line 328, in call_function
    return super().call_function(tx, args, kwargs)
  File "torch/_dynamo/variables/functions.py", line 129, in call_function
    return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 905, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3069, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3195, in inline_call_
    tracer.run()
  File "torch/_dynamo/symbolic_convert.py", line 1054, in run
    while self.step():
  File "torch/_dynamo/symbolic_convert.py", line 964, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 661, in wrapper
    return inner_fn(self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 1738, in CALL_FUNCTION_EX
    self.call_function(fn, argsvars.items, kwargsvars)
  File "torch/_dynamo/symbolic_convert.py", line 899, in call_function
    self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
  File "torch/_dynamo/variables/functions.py", line 389, in call_function
    return super().call_function(tx, args, kwargs)
  File "torch/_dynamo/variables/functions.py", line 328, in call_function
    return super().call_function(tx, args, kwargs)
  File "torch/_dynamo/variables/functions.py", line 129, in call_function
    return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 905, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3069, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3195, in inline_call_
    tracer.run()
  File "torch/_dynamo/symbolic_convert.py", line 1054, in run
    while self.step():
  File "torch/_dynamo/symbolic_convert.py", line 964, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 661, in wrapper
    return inner_fn(self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 1660, in CALL_FUNCTION
    self.call_function(fn, args, {})
  File "torch/_dynamo/symbolic_convert.py", line 899, in call_function
    self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
  File "torch/_dynamo/variables/functions.py", line 389, in call_function
    return super().call_function(tx, args, kwargs)
  File "torch/_dynamo/variables/functions.py", line 328, in call_function
    return super().call_function(tx, args, kwargs)
  File "torch/_dynamo/variables/functions.py", line 129, in call_function
    return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 905, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3069, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3195, in inline_call_
    tracer.run()
  File "torch/_dynamo/symbolic_convert.py", line 1054, in run
    while self.step():
  File "torch/_dynamo/symbolic_convert.py", line 964, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 661, in wrapper
    return inner_fn(self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 1660, in CALL_FUNCTION
    self.call_function(fn, args, {})
  File "torch/_dynamo/symbolic_convert.py", line 899, in call_function
    self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
  File "torch/_dynamo/variables/functions.py", line 129, in call_function
    return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 905, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3069, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3195, in inline_call_
    tracer.run()
  File "torch/_dynamo/symbolic_convert.py", line 1054, in run
    while self.step():
  File "torch/_dynamo/symbolic_convert.py", line 964, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 661, in wrapper
    return inner_fn(self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 1660, in CALL_FUNCTION
    self.call_function(fn, args, {})
  File "torch/_dynamo/symbolic_convert.py", line 899, in call_function
    self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
  File "torch/_dynamo/variables/functions.py", line 328, in call_function
    return super().call_function(tx, args, kwargs)
  File "torch/_dynamo/variables/functions.py", line 129, in call_function
    return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 905, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3069, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3195, in inline_call_
    tracer.run()
  File "torch/_dynamo/symbolic_convert.py", line 1054, in run
    while self.step():
  File "torch/_dynamo/symbolic_convert.py", line 964, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 661, in wrapper
    return inner_fn(self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 1750, in CALL_FUNCTION_KW
    self.call_function(fn, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 899, in call_function
    self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
  File "torch/_dynamo/variables/functions.py", line 389, in call_function
    return super().call_function(tx, args, kwargs)
  File "torch/_dynamo/variables/functions.py", line 328, in call_function
    return super().call_function(tx, args, kwargs)
  File "torch/_dynamo/variables/functions.py", line 129, in call_function
    return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 905, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3069, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3195, in inline_call_
    tracer.run()
  File "torch/_dynamo/symbolic_convert.py", line 1054, in run
    while self.step():
  File "torch/_dynamo/symbolic_convert.py", line 964, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 661, in wrapper
    return inner_fn(self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 1660, in CALL_FUNCTION
    self.call_function(fn, args, {})
  File "torch/_dynamo/symbolic_convert.py", line 899, in call_function
    self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
  File "torch/_dynamo/variables/functions.py", line 328, in call_function
    return super().call_function(tx, args, kwargs)
  File "torch/_dynamo/variables/functions.py", line 129, in call_function
    return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 905, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3069, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3195, in inline_call_
    tracer.run()
  File "torch/_dynamo/symbolic_convert.py", line 1054, in run
    while self.step():
  File "torch/_dynamo/symbolic_convert.py", line 964, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 661, in wrapper
    return inner_fn(self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 1750, in CALL_FUNCTION_KW
    self.call_function(fn, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 899, in call_function
    self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
  File "torch/_dynamo/variables/functions.py", line 389, in call_function
    return super().call_function(tx, args, kwargs)
  File "torch/_dynamo/variables/functions.py", line 328, in call_function
    return super().call_function(tx, args, kwargs)
  File "torch/_dynamo/variables/functions.py", line 129, in call_function
    return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 905, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3069, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3195, in inline_call_
    tracer.run()
  File "torch/_dynamo/symbolic_convert.py", line 1054, in run
    while self.step():
  File "torch/_dynamo/symbolic_convert.py", line 964, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 661, in wrapper
    return inner_fn(self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 1738, in CALL_FUNCTION_EX
    self.call_function(fn, argsvars.items, kwargsvars)
  File "torch/_dynamo/symbolic_convert.py", line 899, in call_function
    self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
  File "torch/_dynamo/variables/user_defined.py", line 600, in call_function
    return super().call_function(tx, args, kwargs)
  File "torch/_dynamo/variables/base.py", line 405, in call_function
    unimplemented(f"call_function {self} {args} {kwargs}")
  File "torch/_dynamo/exc.py", line 317, in unimplemented
    raise Unsupported(msg, case_name=case_name)
torch._dynamo.exc.Unsupported: call_function UserDefinedClassVariable(<class 'tensordict.tensorclass.NonTensorData'>) [] {'data': ConstantVariable(bool: False), '_metadata': ConstantVariable(NoneType: None), '_is_non_tensor': ConstantVariable(bool: True), 'batch_size': SizeVariable(length=0), 'device': ConstantVariable(NoneType: None), 'names': ConstantVariable(NoneType: None)}

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "torch/_dynamo/symbolic_convert.py", line 661, in wrapper
    return inner_fn(self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 1660, in CALL_FUNCTION
    self.call_function(fn, args, {})
  File "torch/_dynamo/symbolic_convert.py", line 899, in call_function
    self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
  File "torch/_dynamo/variables/ctx_manager.py", line 1222, in call_function
    return self.ctx.exit(tx, *args)
  File "torch/_dynamo/variables/ctx_manager.py", line 184, in exit
    unimplemented(
  File "torch/_dynamo/exc.py", line 316, in unimplemented
    raise Unsupported(msg, case_name=case_name) from from_exc
torch._dynamo.exc.Unsupported: Unsupported context manager TensorDict(
...)'s __exit__ function

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/Users/vmoens/Repos/rl/rl/sota-implementations/cql/cql_online.py", line 185, in main
    loss_td = update(sampled_tensordict)
  File "torch/_dynamo/eval_frame.py", line 573, in _fn
    return fn(*args, **kwargs)
  File "torch/_dynamo/convert_frame.py", line 1379, in __call__
    return self._torchdynamo_orig_callable(
  File "torch/_dynamo/convert_frame.py", line 547, in __call__
    return _compile(
  File "torch/_dynamo/convert_frame.py", line 986, in _compile
    guarded_code = compile_inner(code, one_graph, hooks, transform)
  File "torch/_dynamo/convert_frame.py", line 715, in compile_inner
    return _compile_inner(code, one_graph, hooks, transform)
  File "torch/_utils_internal.py", line 95, in wrapper_function
    return function(*args, **kwargs)
  File "torch/_dynamo/convert_frame.py", line 750, in _compile_inner
    out_code = transform_code_object(code, transform)
  File "torch/_dynamo/bytecode_transformation.py", line 1349, in transform_code_object
    transformations(instructions, code_options)
  File "torch/_dynamo/convert_frame.py", line 231, in _fn
    return fn(*args, **kwargs)
  File "torch/_dynamo/convert_frame.py", line 662, in transform
    tracer.run()
  File "torch/_dynamo/symbolic_convert.py", line 2865, in run
    super().run()
  File "torch/_dynamo/symbolic_convert.py", line 1054, in run
    while self.step():
  File "torch/_dynamo/symbolic_convert.py", line 964, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 661, in wrapper
    return inner_fn(self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 1660, in CALL_FUNCTION
    self.call_function(fn, args, {})
  File "torch/_dynamo/symbolic_convert.py", line 899, in call_function
    self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
  File "torch/_dynamo/variables/lazy.py", line 170, in realize_and_forward
    return getattr(self.realize(), name)(*args, **kwargs)
  File "torch/_dynamo/variables/nn_module.py", line 914, in call_function
    return variables.UserFunctionVariable(fn, source=source).call_function(
  File "torch/_dynamo/variables/functions.py", line 328, in call_function
    return super().call_function(tx, args, kwargs)
  File "torch/_dynamo/variables/functions.py", line 129, in call_function
    return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 905, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3069, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3195, in inline_call_
    tracer.run()
  File "torch/_dynamo/symbolic_convert.py", line 1054, in run
    while self.step():
  File "torch/_dynamo/symbolic_convert.py", line 964, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 661, in wrapper
    return inner_fn(self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 1660, in CALL_FUNCTION
    self.call_function(fn, args, {})
  File "torch/_dynamo/symbolic_convert.py", line 899, in call_function
    self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
  File "torch/_dynamo/variables/functions.py", line 129, in call_function
    return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 905, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3069, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3195, in inline_call_
    tracer.run()
  File "torch/_dynamo/symbolic_convert.py", line 1054, in run
    while self.step():
  File "torch/_dynamo/symbolic_convert.py", line 964, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 661, in wrapper
    return inner_fn(self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 1738, in CALL_FUNCTION_EX
    self.call_function(fn, argsvars.items, kwargsvars)
  File "torch/_dynamo/symbolic_convert.py", line 899, in call_function
    self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
  File "torch/_dynamo/variables/functions.py", line 389, in call_function
    return super().call_function(tx, args, kwargs)
  File "torch/_dynamo/variables/functions.py", line 328, in call_function
    return super().call_function(tx, args, kwargs)
  File "torch/_dynamo/variables/functions.py", line 129, in call_function
    return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 905, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3069, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3195, in inline_call_
    tracer.run()
  File "torch/_dynamo/symbolic_convert.py", line 1054, in run
    while self.step():
  File "torch/_dynamo/symbolic_convert.py", line 964, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 661, in wrapper
    return inner_fn(self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 1738, in CALL_FUNCTION_EX
    self.call_function(fn, argsvars.items, kwargsvars)
  File "torch/_dynamo/symbolic_convert.py", line 899, in call_function
    self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
  File "torch/_dynamo/variables/functions.py", line 328, in call_function
    return super().call_function(tx, args, kwargs)
  File "torch/_dynamo/variables/functions.py", line 129, in call_function
    return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 905, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3069, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3195, in inline_call_
    tracer.run()
  File "torch/_dynamo/symbolic_convert.py", line 1054, in run
    while self.step():
  File "torch/_dynamo/symbolic_convert.py", line 964, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 661, in wrapper
    return inner_fn(self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 1738, in CALL_FUNCTION_EX
    self.call_function(fn, argsvars.items, kwargsvars)
  File "torch/_dynamo/symbolic_convert.py", line 899, in call_function
    self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
  File "torch/_dynamo/variables/functions.py", line 328, in call_function
    return super().call_function(tx, args, kwargs)
  File "torch/_dynamo/variables/functions.py", line 129, in call_function
    return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 905, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3069, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3195, in inline_call_
    tracer.run()
  File "torch/_dynamo/symbolic_convert.py", line 1054, in run
    while self.step():
  File "torch/_dynamo/symbolic_convert.py", line 964, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 661, in wrapper
    return inner_fn(self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 1660, in CALL_FUNCTION
    self.call_function(fn, args, {})
  File "torch/_dynamo/symbolic_convert.py", line 899, in call_function
    self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
  File "torch/_dynamo/variables/functions.py", line 389, in call_function
    return super().call_function(tx, args, kwargs)
  File "torch/_dynamo/variables/functions.py", line 328, in call_function
    return super().call_function(tx, args, kwargs)
  File "torch/_dynamo/variables/functions.py", line 129, in call_function
    return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 905, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3069, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3195, in inline_call_
    tracer.run()
  File "torch/_dynamo/symbolic_convert.py", line 1054, in run
    while self.step():
  File "torch/_dynamo/symbolic_convert.py", line 964, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 661, in wrapper
    return inner_fn(self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 1738, in CALL_FUNCTION_EX
    self.call_function(fn, argsvars.items, kwargsvars)
  File "torch/_dynamo/symbolic_convert.py", line 899, in call_function
    self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
  File "torch/_dynamo/variables/functions.py", line 328, in call_function
    return super().call_function(tx, args, kwargs)
  File "torch/_dynamo/variables/functions.py", line 129, in call_function
    return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 905, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3069, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3195, in inline_call_
    tracer.run()
  File "torch/_dynamo/symbolic_convert.py", line 1054, in run
    while self.step():
  File "torch/_dynamo/symbolic_convert.py", line 964, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 661, in wrapper
    return inner_fn(self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 1660, in CALL_FUNCTION
    self.call_function(fn, args, {})
  File "torch/_dynamo/symbolic_convert.py", line 899, in call_function
    self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
  File "torch/_dynamo/variables/functions.py", line 389, in call_function
    return super().call_function(tx, args, kwargs)
  File "torch/_dynamo/variables/functions.py", line 328, in call_function
    return super().call_function(tx, args, kwargs)
  File "torch/_dynamo/variables/functions.py", line 129, in call_function
    return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 905, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3069, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "torch/_dynamo/symbolic_convert.py", line 3195, in inline_call_
    tracer.run()
  File "torch/_dynamo/symbolic_convert.py", line 1054, in run
    while self.step():
  File "torch/_dynamo/symbolic_convert.py", line 964, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "torch/_dynamo/symbolic_convert.py", line 667, in wrapper
    unimplemented("Graph break under GenericContextWrappingVariable")
  File "torch/_dynamo/exc.py", line 317, in unimplemented
    raise Unsupported(msg, case_name=case_name)
torch._dynamo.exc.Unsupported: Graph break under GenericContextWrappingVariable

from user code:
   File "/Users/vmoens/Repos/rl/rl/sota-implementations/cql/cql_online.py", line 121, in update
    loss_td = loss_module(sampled_tensordict)
  File "torch/nn/modules/module.py", line 1844, in _call_impl
    return inner()
  File "torch/nn/modules/module.py", line 1793, in inner
    result = forward_call(*args, **kwargs)
  File "torchrl/objectives/common.py", line 55, in new_forward
    return func(self, *args, **kwargs)
  File "/Users/vmoens/Repos/rl/tensordict/tensordict/nn/common.py", line 316, in wrapper
    return func(_self, tensordict, *args, **kwargs)
  File "torchrl/objectives/cql.py", line 518, in forward
    q_loss, metadata = self.q_loss(tensordict)
  File "torchrl/objectives/common.py", line 55, in new_forward
    return func(self, *args, **kwargs)
  File "torchrl/objectives/cql.py", line 686, in q_loss
    target_value = self._get_value_v(
  File "torchrl/objectives/cql.py", line 631, in _get_value_v
    with set_exploration_type(ExplorationType.RANDOM), actor_params.data.to_module(


You can suppress this exception and fall back to eager by setting:
    import torch._dynamo
    torch._dynamo.config.suppress_errors = True
```

</details>


### Versions

nightlies

cc @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @amjames