Repository: pytorch
Issue: [TensorDict - compile] Dynamo doens't like a simple class decorator - but a function is fine
Labels: triaged,oncall: pt2,module: dynamo,dynamo-triage-june2024,dynamo-functools

### 🐛 Describe the bug

In tensordict we have this decorator:
```python
class as_decorator:
    """Converts a method to a decorator.

    Examples:
        >>> from tensordict import TensorDict
        >>> data = TensorDict({}, [])
        >>> with data.lock_(): # lock_ is decorated
        ...     assert data.is_locked
        >>> assert not data.is_locked
    """

    def __init__(self, attr=None):
        self.attr = attr

    def __call__(self, func):
        if self.attr is not None:

            @wraps(func)
            def new_func(_self, *args, **kwargs):
                _attr_pre = getattr(_self, self.attr)
                out = func(_self, *args, **kwargs)
                _attr_post = getattr(_self, self.attr)
                if out is not None:
                    if _attr_post is not _attr_pre:
                        out._last_op = (new_func.__name__, (args, kwargs, _self))
                    else:
                        out._last_op = None
                return out

        else:

            @wraps(func)
            def new_func(_self, *args, **kwargs):
                out = func(_self, *args, **kwargs)
                if out is not None:
                    out._last_op = (new_func.__name__, (args, kwargs, _self))
                return out

        return new_func
class Foo:
    @as_decorator("is_locked)
    def lock_(self):
        ...
```
Dynamo doesn't like that:
```
  File "/Users/vmoens/venv/rl/lib/python3.11/site-packages/torch/_dynamo/variables/builder.py", line 2504, in create
    unimplemented(
  File "/Users/vmoens/venv/rl/lib/python3.11/site-packages/torch/_dynamo/exc.py", line 221, in unimplemented
    raise Unsupported(msg)
torch._dynamo.exc.Unsupported: Unexpected type in sourceless builder tensordict.utils.as_decorator

from user code:
   File "/Users/vmoens/Repos/rl/tensordict/test/test_compile.py", line 216, in locked_op
    td2 = td + 1
  File "/Users/vmoens/Repos/rl/tensordict/tensordict/base.py", line 6577, in __add__
    return self.add(other)
  File "/Users/vmoens/Repos/rl/tensordict/tensordict/base.py", line 7104, in add
    return self._fast_apply(
  File "/Users/vmoens/Repos/rl/tensordict/tensordict/base.py", line 6102, in _fast_apply
    result.lock_()

Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information
```

However rewriting that as a function works ok
```python
def as_decorator(attr=None):
    """Converts a method to a decorator.

    Examples:
        >>> from tensordict import TensorDict
        >>> data = TensorDict({}, [])
        >>> with data.lock_(): # lock_ is decorated
        ...     assert data.is_locked
        >>> assert not data.is_locked
    """

    def __call__(func):
        if attr is not None:

            @wraps(func)
            def new_func(_self, *args, **kwargs):
                _attr_pre = getattr(_self, attr)
                out = func(_self, *args, **kwargs)
                _attr_post = getattr(_self, attr)
                if out is not None:
                    if _attr_post is not _attr_pre:
                        out._last_op = (new_func.__name__, (args, kwargs, _self))
                    else:
                        out._last_op = None
                return out

        else:

            @wraps(func)
            def new_func(_self, *args, **kwargs):
                out = func(_self, *args, **kwargs)
                if out is not None:
                    out._last_op = (new_func.__name__, (args, kwargs, _self))
                return out

        return new_func
    return __call__
class Foo:
    @as_decorator("is_locked)
    def lock_(self):
        ...
```

### Versions

'2.5.0.dev20240709'

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