Since torch functions often allow for models to either be from disk as fpath or from a loaded object, we provide a similar option here. Arguments can either be a path on disk or from model itself. Args: model (Union[str, torch.Model]): Location of the model as f
(
self,
model: Union[str, Module],
default_converter: ModelFileConverter = None,
network_metadata: NetworkMetadata = None,
)
| 256 | |
| 257 | class TorchModelFile(NNModelFile): |
| 258 | def __init__( |
| 259 | self, |
| 260 | model: Union[str, Module], |
| 261 | default_converter: ModelFileConverter = None, |
| 262 | network_metadata: NetworkMetadata = None, |
| 263 | ): |
| 264 | """ |
| 265 | Since torch functions often allow for models to either be from disk as fpath or from a loaded object, |
| 266 | we provide a similar option here. Arguments can either be a path on disk or from model itself. |
| 267 | |
| 268 | Args: |
| 269 | model (Union[str, torch.Model]): Location of the model as fpath OR loaded torch.Model object. |
| 270 | """ |
| 271 | super().__init__(default_converter, network_metadata) |
| 272 | |
| 273 | if isinstance(model, Module): |
| 274 | self.is_loaded = True |
| 275 | self.fpath = None |
| 276 | self.model = model |
| 277 | else: |
| 278 | self.is_loaded = False |
| 279 | self.fpath = model |
| 280 | self.model = None |
| 281 | |
| 282 | def load_model(self) -> Module: |
| 283 | """ |