ray.train.torch.TorchCheckpoint.get_model

TorchCheckpoint.get_model(model: Optional[torch.nn.modules.module.Module] = None) → torch.nn.modules.module.Module

Retrieve the model stored in this checkpoint.

Parameters

model – If the checkpoint contains a model state dict, and not the model itself, then the state dict will be loaded to this model. Otherwise, the model will be discarded.