nvflare.app_opt.pt.job_config.model module

class PTModel(model, persistor: ModelPersistor | None = None, locator: ModelLocator | None = None)[source]

Bases: object

PyTorch model wrapper.

If model is an nn.Module, add a PTFileModelPersistor with the model and a TFModelPersistor.

Parameters:

  • model (any) – model

  • persistor (optional, ModelPersistor) – how to persistor the model.

  • locator (optional, ModelLocator) – how to locate the model.

add_to_fed_job(job, ctx)[source]

This method is used by Job API.

Parameters:

  • job – the Job object to add to

  • ctx – Job Context

Returns:

dictionary of ids of component added