Files
pytorch/test/mkl_verbose.py
Yuanhao Ji 706f7d1f22 Enable UFMT on test/jit_hooks, test/lazy and some files (#123807)
Part of: #123062

Ran lintrunner on:

- `test/jit_hooks`
- `test/lazy`
- `test/linear.py`
- `test/load_torchscript_model.py`
- `test/mkl_verbose.py`
- `test/mkldnn_verbose.py`

with command:

```bash
lintrunner -a --take UFMT --all-files
```

Co-authored-by: Edward Z. Yang <ezyang@fb.com>
Pull Request resolved: https://github.com/pytorch/pytorch/pull/123807
Approved by: https://github.com/ezyang
2024-04-12 03:39:38 +00:00

21 lines
434 B
Python

import argparse
import torch
def run_model(level):
m = torch.nn.Linear(20, 30)
input = torch.randn(128, 20)
with torch.backends.mkl.verbose(level):
m(input)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--verbose-level", default=0, type=int)
args = parser.parse_args()
try:
run_model(args.verbose_level)
except Exception as e:
print(e)