Reformated onnx_transforms.py for ruff check

Signed-off-by: abhishek-singh591 <[email protected]>

Author: abhishek-singh591

QEfficient/base/onnx_transforms.py

@@ -3,16 +3,18 @@
 # Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries.
 # SPDX-License-Identifier: BSD-3-Clause
 #
-# ----------------------------------------------------------------------------
+# -----------------------------------------------------------------------------
+
 import os
 from typing import Optional, Tuple
 import numpy as np
 from onnx import ModelProto, external_data_helper, numpy_helper
 from concurrent.futures import ThreadPoolExecutor
 
+
 class OnnxTransform:
     """
-    OnnxTransform is the base class for graph modifications on exported onnx.
+    OnnxTransform is the base class for graph modifications on exported ONNX.
     """
 
     def __init__(self):
@@ -22,20 +24,22 @@ def __init__(self):
     def apply(cls, model: ModelProto, **kwargs) -> Tuple[ModelProto, bool]:
         """
         Override this class to apply a transformation.
+
         :param model: The model's ONNX graph to transform
         :param kwargs: Parameters needed for specific transforms. All transforms should take **kwargs to ignore unneeded kwargs.
-
         :returns: ONNX graph after applying the transform
         :returns: Boolean indicating whether transform was applied
         """
         raise NotImplementedError("Use subclasses for ONNX transform")
-    
+
+
 class FP16ClipTransform(OnnxTransform):
     @classmethod
     def apply(cls, model: ModelProto, *, onnx_base_dir: Optional[str] = None, **kwargs) -> Tuple[ModelProto, bool]:
         finfo = np.finfo(np.float16)
         fp16_max = finfo.max
         fp16_min = finfo.min
+
         def clip_tensor(tensor):
             nptensor = numpy_helper.to_array(tensor, onnx_base_dir)
             if nptensor.dtype == np.float32 and (np.any(nptensor > fp16_max) or np.any(nptensor < fp16_min)):
@@ -51,6 +55,7 @@ def clip_tensor(tensor):
         transformed = any(results)
         return model, transformed
 
+
 class SplitTensorsTransform(OnnxTransform):
     @classmethod
     def apply(
@@ -63,14 +68,13 @@ def apply(
         size_threshold: int = 1024,
         **kwargs,
     ) -> Tuple[ModelProto, bool]:
-        
         external_data_helper.load_external_data_for_model(model, onnx_base_dir)
         tensors = external_data_helper._get_all_tensors(model)
         file_assignments = []
         file_num = 0
         current_file_size = 0
-        transformed = False 
-        
+        transformed = False
+
         for tensor in tensors:
             if tensor.HasField("raw_data") and (tsize := len(tensor.raw_data)) > size_threshold:
                 transformed = True
@@ -85,5 +89,5 @@ def process_tensor(args):
             external_data_helper.set_external_data(tensor, file_name)
 
         with ThreadPoolExecutor(max_workers=os.cpu_count() * 4) as executor:
-            executor.map(process_tensor, file_assignments)
+            list(executor.map(process_tensor, file_assignments))
         return model, transformed