keras/src/utils/torch_utils.py/TorchModuleWrapper/__init__
class TorchModuleWrapper:    def __init__(self, module, name=None, **kwargs):
        super().__init__(name=name, **kwargs)
        import torch.nn as nn

        from keras.src.backend.torch.core import get_device

        if (
            isinstance(module, nn.modules.lazy.LazyModuleMixin)
            and module.has_uninitialized_params()
        ):
            raise ValueError(
                "LazyModules are not supported unless they "
                "are already initialized. "
                f"Received uninitialized LazyModule: module={module}"
            )

        self.module = module.to(get_device())
        self._track_module_parameters()