class maze.core.agent.torch_model.TorchModel(device: str)

Base class for any torch model.


device – Device the networks should be located on (cpu or cuda)

property device

Returns the device the networks are located on.

abstract eval()None

Set all networks to eval mode.

abstract load_state_dict(state_dict: Dict)None

Set state dict of all encapsulated networks. :param state_dict: The torch state dictionary.

property num_params

Returns overall number of network parameters.

abstract parameters() → List[torch.Tensor]

Returns all parameters of all networks.

abstract state_dict() → Dict

Return state dict composed of state dicts of all encapsulated networks.

abstract to(device: str)None

Move all networks to the specified device. :param device: The target device.

abstract train()None

Set all networks to training mode.