nip.utils.torch.TensorDictCloneKeys#

class nip.utils.torch.TensorDictCloneKeys(*args, **kwargs)[source]#

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, ...] | Iterable[str | Tuple[str, ...]], out_keys: str | Tuple[str, ...] | Iterable[str | Tuple[str, ...]])[source]#

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

forward(tensordict: TensorDictBase) TensorDictBase[source]#

Clone the keys of the input TensorDict.

Parameters:

tensordict (TensorDictBase) – The input TensorDict.

Returns:

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