Skip to content

dict_of_tensor_mixin

dict_of_tensor_mixin

Mixin for modules that store and load nested ParameterDict state.

DictOfTensorMixin 

DictOfTensorMixin(params_dict=None)

Bases: ModuleAttrMixin

Store nested tensors under params_dict and rebuild them on load.

Source code in src/versatil/common/dict_of_tensor_mixin.py 
14
15
16
17
18
def __init__(self, params_dict: nn.ParameterDict | None = None) -> None:
    super().__init__()
    self.params_dict = (
        params_dict if params_dict is not None else nn.ParameterDict()
    )