Arm backend: Use _pytree tree_flatten in model evaluator (#11133)

### Summary
Updates model evaluator to use tree_flatten for output tensors when
checking int8 deviation.
* Removes in script function which flattens the tensors.

Author: SaoirseARM

backends/arm/util/arm_model_evaluator.py

@@ -18,6 +18,7 @@
 
 import torch
 from torch.nn.modules import Module
+from torch.utils._pytree import tree_flatten
 from torch.utils.data import DataLoader
 from torchvision import datasets, transforms  # type: ignore[import-untyped]
 
@@ -28,20 +29,6 @@
 logger.setLevel(logging.INFO)
 
 
-def flatten_args(args) -> tuple | list:
-    flattened_args: list = []
-    if isinstance(args, torch.Tensor):
-        return [args]
-
-    for arg in args:
-        if isinstance(arg, (tuple, list)):
-            flattened_args.extend(arg)
-        else:
-            flattened_args.append(arg)
-
-    return tuple(flattened_args)
-
-
 class GenericModelEvaluator:
     REQUIRES_CONFIG = False
 
@@ -72,8 +59,8 @@ def get_model_error(self) -> defaultdict:
         - Maximum percentage error
         - Mean absolute error
         """
-        fp32_outputs = flatten_args(self.fp32_model(*self.example_input))
-        int8_outputs = flatten_args(self.int8_model(*self.example_input))
+        fp32_outputs, _ = tree_flatten(self.fp32_model(*self.example_input))
+        int8_outputs, _ = tree_flatten(self.int8_model(*self.example_input))
 
         model_error_dict = defaultdict(list)