Addressed PR comments

Author: suryasidd

examples/openvino/aot_optimize_and_infer.py

@@ -108,7 +108,6 @@ def infer_model(
     inputs,
     num_iter: int,
     warmup_iter: int,
-    input_path: str,
     output_path: str,
 ) -> float:
     """
@@ -118,34 +117,29 @@ def infer_model(
     :param inputs: The inputs for the model.
     :param num_iter: The number of iterations to execute inference for timing.
     :param warmup_iter: The number of iterations to execute inference for warmup before timing.
-    :param input_path: Path to the input tensor file to read the input for inference.
     :param output_path: Path to the output tensor file to save the output of inference..
     :return: The average inference timing.
     """
-    # 1: Load model from buffer
+    # Load model from buffer
     executorch_module = _load_for_executorch_from_buffer(exec_prog.buffer)
 
-    # 2: Initialize inputs
-    if input_path:
-        inputs = (torch.load(input_path, weights_only=False),)
-
-    # 3: Execute warmup
+    # Execute warmup
     for _i in range(warmup_iter):
         out = executorch_module.run_method("forward", inputs)
 
-    # 4: Execute inference and measure timing
+    # Execute inference and measure timing
     time_total = 0.0
     for _i in range(num_iter):
         time_start = time.time()
         out = executorch_module.run_method("forward", inputs)
         time_end = time.time()
         time_total += time_end - time_start
 
-    # 5: Save output tensor as raw tensor file
+    # Save output tensor as raw tensor file
     if output_path:
         torch.save(out, output_path)
 
-    # 6: Return average inference timing
+    # Return average inference timing
     return time_total / float(num_iter)
 
 
@@ -159,10 +153,10 @@ def validate_model(
     :param calibration_dataset: A DataLoader containing calibration data.
     :return: The accuracy score of the model.
     """
-    # 1: Load model from buffer
+    # Load model from buffer
     executorch_module = _load_for_executorch_from_buffer(exec_prog.buffer)
 
-    # 2: Iterate over the dataset and run the executor
+    # Iterate over the dataset and run the executor
     predictions = []
     targets = []
     for _idx, data in enumerate(calibration_dataset):
@@ -171,7 +165,7 @@ def validate_model(
         out = executorch_module.run_method("forward", (feature,))
         predictions.extend(torch.stack(out).reshape(-1, 1000).argmax(-1))
 
-    # 1: Check accuracy
+    # Check accuracy
     return accuracy_score(predictions, targets)
 
 
@@ -230,12 +224,14 @@ def main(  # noqa: C901
         msg = "Input shape must be a list or tuple."
         raise ValueError(msg)
     # Provide input
-    if suite == "huggingface":
-        if hasattr(model, 'config') and hasattr(model.config, 'vocab_size'):
+    if input_path:
+        example_args = (torch.load(input_path, weights_only=False),)
+    elif suite == "huggingface":
+        if hasattr(model, "config") and hasattr(model.config, "vocab_size"):
             vocab_size = model.config.vocab_size
         else:
             vocab_size = 30522
-        example_args = (torch.randint(0, vocab_size, input_shape, dtype=torch.int64), )
+        example_args = (torch.randint(0, vocab_size, input_shape, dtype=torch.int64),)
     else:
         example_args = (torch.randn(*input_shape),)
 
@@ -306,7 +302,7 @@ def transform_fn(x):
     if infer:
         print("Start inference of the model:")
         avg_time = infer_model(
-            exec_prog, example_args, num_iter, warmup_iter, input_path, output_path
+            exec_prog, example_args, num_iter, warmup_iter, output_path
         )
         print(f"Average inference time: {avg_time}")