nip.utils.torch.TensorDictPrint#
- class nip.utils.torch.TensorDictPrint(*args, **kwargs)[source]#
Print information about an input tensordict.
- Parameters:
Methods Summary
__init__(keys[, name, print_nan_proportion])Initialize internal Module state, shared by both nn.Module and ScriptModule.
forward(tensordict)Print the information about the tensors in the input tensordict.
Attributes
T_destinationcall_super_initdump_patchesin_keysout_keysout_keys_sourcetrainingMethods
- __init__(keys: str | tuple[str | tuple[NestedKey, ...], ...] | Iterable[str | tuple[str | tuple[NestedKey, ...], ...]], name: str | None = None, print_nan_proportion: bool = False)[source]#
Initialize internal Module state, shared by both nn.Module and ScriptModule.
- forward(tensordict: TensorDictBase) TensorDictBase[source]#
Print the information about the tensors in the input tensordict.
- Parameters:
tensordict (TensorDictBase) – The input tensordict.
- Returns:
tensordict (TensorDictBase) – The input tensordict, unchanged.