1import argparse 2 3import torch 4 5 6class Module(torch.nn.Module): 7 def __init__(self) -> None: 8 super().__init__() 9 self.conv = torch.nn.Conv2d(1, 10, 5, 1) 10 11 def forward(self, x): 12 y = self.conv(x) 13 return y 14 15 16def run_model(level): 17 m = Module().eval() 18 d = torch.rand(1, 1, 112, 112) 19 with torch.backends.mkldnn.verbose(level): 20 m(d) 21 22 23if __name__ == "__main__": 24 parser = argparse.ArgumentParser() 25 parser.add_argument("--verbose-level", default=0, type=int) 26 args = parser.parse_args() 27 try: 28 run_model(args.verbose_level) 29 except Exception as e: 30 print(e) 31