code cleanup

Signed-off-by: Gal Hubara Agam <96368689+galagam@users.noreply.github.com>

Author: galagam

modelopt/onnx/autocast/referencerunner.py

@@ -121,54 +121,34 @@ def _load_inputs(self, inputs):
         return data_loader
 
     def _get_ort_runner(self, model):
-        import onnxruntime as ort
         from polygraphy.backend.onnx import BytesFromOnnx
         from polygraphy.backend.onnxrt import OnnxrtRunner, SessionFromOnnx
 
         # Check if model has external data by checking:
         # 1. If any initializer has data_location set to EXTERNAL (even if data is loaded)
         # 2. If model size would exceed 2GB (indicating need for external data)
-        has_external_data = onnx_utils.check_model_uses_external_data(self.model)
-
-        # Also check if model would be too large (>2GB) for SerializeToString
-        # This handles cases where model was loaded with external data already loaded
-        if not has_external_data:
-            try:
-                # Try to estimate size by serializing the model
-                # If it fails or exceeds 2GB, we need file-based approach
-                model_size = len(self.model.SerializeToString())
-                if model_size > 2 * (1024**3):  # 2GB threshold
-                    has_external_data = True
-                    logger.debug(
-                        f"Model size ({model_size / (1024**3):.2f} GB) exceeds 2GB, using file-based approach"
-                    )
-            except (ValueError, AttributeError) as e:
-                # SerializeToString failed (likely >2GB limit), use file-based approach
-                if "exceeds maximum protobuf size" in str(e) or "2GB" in str(e):
-                    has_external_data = True
-                    logger.debug("Model exceeds protobuf 2GB limit, using file-based approach")
-
-        if has_external_data:
+        needs_external_data = onnx_utils.check_model_uses_external_data(
+            self.model
+        ) or self.model.ByteSize() > 2 * (1024**3)
+        if needs_external_data:
             logger.debug("Model has external data, using file-based approach")
             # Get the actual ONNX ModelProto from ModifyOutputs wrapper
             modified_model = model()
 
-            # Use a persistent temp file to handle external data files properly
+            # Use a persistent temp file, because we need the file to be present in an broader context
             tmp_file = tempfile.NamedTemporaryFile(suffix=".onnx", delete=False)
             tmp_file.close()
             tmp_file_path = tmp_file.name
             onnx_utils.save_onnx(modified_model, tmp_file_path, save_as_external_data=True)
             logger.debug(f"Model with all outputs saved to {tmp_file_path}")
-            session = ort.InferenceSession(tmp_file_path, providers=self.providers)
-            runners = [OnnxrtRunner(lambda: session)]
+            build_onnxrt_session = SessionFromOnnx(tmp_file_path, providers=self.providers)
 
         else:
             # For models without external data, use the original BytesFromOnnx approach (no tmp files)
             logger.debug("Model has no external data, using BytesFromOnnx approach")
             serialize_onnx = BytesFromOnnx(model)
             build_onnxrt_session = SessionFromOnnx(serialize_onnx, providers=self.providers)
-            runners = [OnnxrtRunner(build_onnxrt_session)]
-
+        runners = [OnnxrtRunner(build_onnxrt_session)]
         return runners
 
     def run(self, inputs=None):