as_tensordict_module

class tensordict.nn.as_tensordict_module(*, in_keys: Union[List[NestedKey], NestedKey], out_keys: Union[List[NestedKey], NestedKey])

將函式轉換為 TensorDictModule 的裝飾器。

引數:

• in_keys (List[NestedKey] | NestedKey | None, optional) – 結果 TensorDictModule 的輸入鍵。

• out_keys (List[NestedKey] | NestedKey | None, optional) – 結果 TensorDictModule 的輸出鍵。

返回:

可以應用於函式的裝飾器，將其轉換為 TensorDictModule。

返回型別:

Callable

示例

>>> class MyClass:
...     @as_tensordict_module(in_keys="c", out_keys="d")
...     def my_method(self, c):
...         return c + 1
>>> obj = MyClass()
>>> result = obj.my_method(TensorDict(c=0))
>>> print(result["d"])  # prints: 1
>>> @as_tensordict_module(in_keys="c", out_keys="d")
... def my_function(c):
...     return c + 1
>>> result = my_function(TensorDict(c=0))
>>> print(result["d"])  # prints: 1