def test_unwrap_lightning_module():
model = BoringModel()
wrapped_model = _LightningPrecisionModuleWrapperBase(model)
wrapped_model = _LightningModuleWrapperBase(wrapped_model)
wrapped_model = DataParallel(wrapped_model)
assert unwrap_lightning_module(wrapped_model) == model
def unwrap_lightning_module_sharded(
wrapped_model: nn.Module) -> "pl.LightningModule":
model = wrapped_model
if isinstance(model, ShardedDataParallel):
model = model.module
return unwrap_lightning_module(model)
def lightning_module(self):
return unwrap_lightning_module(self._model)
def lightning_module(self) -> LightningModule:
"""Returns the pure LightningModule without potential wrappers"""
return unwrap_lightning_module(self._model)
def lightning_module(self) -> "pl.LightningModule":
model = self._model
if isinstance(model, BaguaDistributedDataParallel):
model = model.module
return unwrap_lightning_module(model) # type: ignore[arg-type]
def lightning_module(self) -> Optional["pl.LightningModule"]:
return unwrap_lightning_module(self.model) if self.model is not None else None
def lightning_module(self) -> Optional["pl.LightningModule"]:
"""Returns the pure LightningModule without potential wrappers."""
return unwrap_lightning_module(self.model) if self.model is not None else None
def lightning_module(self) -> Optional["pl.LightningModule"]:
model = self.model
if isinstance(model, BaguaDistributedDataParallel):
model = model.module
return unwrap_lightning_module(model) if model is not None else None