nip.utils.torch.TensorDictCloneKeys

class nip.utils.torch.TensorDictCloneKeys(*args, **kwargs)

Clone the keys of a TensorDict.

Parameters:

• in_keys (NestedKey | Iterable[NestedKey]) – The keys to clone.

• out_keys (NestedKey | Iterable[NestedKey]) – The keys to store the cloned values in.

Methods Summary

__init__(in_keys, out_keys)	Initialize internal Module state, shared by both nn.Module and ScriptModule.
forward(tensordict)	Clone the keys of the input TensorDict.

Attributes

T_destination
call_super_init
dump_patches
in_keys
out_keys
out_keys_source
training

Methods

__init__(in_keys: str | tuple[str | tuple[NestedKey, ...], ...] | Iterable[str | tuple[str | tuple[NestedKey, ...], ...]], out_keys: str | tuple[str | tuple[NestedKey, ...], ...] | Iterable[str | tuple[str | tuple[NestedKey, ...], ...]])

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(tensordict: TensorDictBase) → TensorDictBase

Clone the keys of the input TensorDict.

Parameters:

tensordict (TensorDictBase) – The input TensorDict.

Returns:

cloned_tensordict (TensorDictBase) – The input TensorDict with the keys cloned.