gpytorchwrapper.src.models.model_save

Functions

save_model(state_dict, config, ...)

Save the model to a .pth file
gpytorchwrapper.src.models.model_save.save_model(state_dict: dict, config: dict, training_data: dict, metrics: dict, output: str, directory: Path) → None[source]

Save the model to a .pth file

Parameters:

  • state_dict (dict) – The state of the model

  • config (dict) – The configuration dictionary to setup the training

  • training_data (dict) – Contains the training inputs and outputs used

  • metrics (dict) – Dictionary containing the metrics on the training and test sets

  • output (str) – The name of the output file

  • directory (pathlib.Path) – The output directory

Return type:

None