| 2 | |
| 3 | |
| 4 | class BaseModel(torch.nn.Module): |
| 5 | def load(self, path): |
| 6 | """Load model from file. |
| 7 | |
| 8 | Args: |
| 9 | path (str): file path |
| 10 | """ |
| 11 | parameters = torch.load(path, map_location=torch.device('cpu')) |
| 12 | |
| 13 | if "optimizer" in parameters: |
| 14 | parameters = parameters["model"] |
| 15 | |
| 16 | self.load_state_dict(parameters) |
nothing calls this directly
no outgoing calls
no test coverage detected