module_attr_mixin

module_attr_mixin

Mixin for nn.Module device and dtype introspection.

ModuleAttrMixin

ModuleAttrMixin()

Bases: Module

Expose module-level device and dtype properties.

A non-persistent reference buffer lets stateless and lazily initialized modules track .to(...) calls without adding trainable parameters or polluting checkpoints.

Source code in src/versatil/common/module_attr_mixin.py

def __init__(self) -> None:
    super().__init__()
    self.register_buffer(
        "_module_attr_reference",
        torch.empty(0),
        persistent=False,
    )

device property writable

device

Return the module's current device.

dtype property writable

dtype

Return the module's current floating-point dtype.