Add backend parameter to the sample

sovrasov · Jan 21, 2024 · ea6a4ca · ea6a4ca
1 parent 744e746
commit ea6a4ca
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/samples/classification.py b/samples/classification.py
@@ -27,6 +27,8 @@
                         help='Device to store the model.')
     parser.add_argument('--model', choices=list(pt_models.keys()),
                         type=str, default='resnet18')
+    parser.add_argument('--backend', choices=list(['pytorch', 'aten']),
+                        type=str, default='pytorch')
     parser.add_argument('--result', type=str, default=None)
     args = parser.parse_args()
 
@@ -42,6 +44,7 @@
 
     macs, params = get_model_complexity_info(net, (3, 224, 224),
                                              as_strings=True,
+                                             backend=args.backend,
                                              print_per_layer_stat=True,
                                              ost=ost)
     print('{:<30}  {:<8}'.format('Computational complexity: ', macs))