(model, **compile_options)
| 36 | |
| 37 | |
| 38 | def compile_model(model, **compile_options): |
| 39 | # Compile the model with torch 2.0 |
| 40 | if hasattr(model, 'compile'): |
| 41 | model = model.compile(**compile_options) |
| 42 | elif version.parse(torch.__version__) >= version.parse('2.0.0.dev'): |
| 43 | model = torch.compile(model, **compile_options) |
| 44 | else: |
| 45 | print( |
| 46 | 'Compiling model needs torch version > 2.0.0, ' |
| 47 | f'your torch version is: {torch.__version__}, origin model will be returned.' |
| 48 | ) |
| 49 | return model |
| 50 | |
| 51 | |
| 52 | def init_dist(launcher: str, backend: str = 'nccl', **kwargs) -> None: |
no test coverage detected
searching dependent graphs…