Load the pre-trained models.
()
| 16 | |
| 17 | |
| 18 | def load_models(): |
| 19 | """Load the pre-trained models.""" |
| 20 | models = {} |
| 21 | for style in styles.values(): |
| 22 | model = Transformer() |
| 23 | model_file_path = os.path.join(model_path, style + "_net_G_float.pth") |
| 24 | |
| 25 | logging.info(f"Loading { style } from { model_file_path }...") |
| 26 | model.load_state_dict( |
| 27 | torch.load(os.path.join(model_path, style + "_net_G_float.pth")) |
| 28 | ) |
| 29 | model.eval() |
| 30 | if gpu: |
| 31 | model.cuda() |
| 32 | else: |
| 33 | model.float() |
| 34 | |
| 35 | models[style] = model |
| 36 | |
| 37 | return models |
| 38 | |
| 39 | |
| 40 | models = load_models() |
no test coverage detected