Enhance type hinting in MonaiBundleInferenceOperator

- Added TYPE_CHECKING imports for torch to improve type hinting and static analysis.
- Utilized type casting for model loading return values to ensure proper type safety and clarity.

Signed-off-by: Victor Chang <[email protected]>

Author: mocsharp

monai/deploy/operators/monai_bundle_inference_operator.py

@@ -20,7 +20,7 @@
 from copy import deepcopy
 from pathlib import Path
 from threading import Lock
-from typing import Any, Dict, List, Optional, Tuple, Type, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type, Union, cast
 
 import numpy as np
 
@@ -29,6 +29,9 @@
 
 from .inference_operator import InferenceOperator
 
+if TYPE_CHECKING:
+    import torch
+
 MONAI_UTILS = "monai.utils"
 nibabel, _ = optional_import("nibabel", "3.2.1")
 torch, _ = optional_import("torch", "1.10.2")
@@ -93,12 +96,12 @@ def _load_model_from_directory_bundle(bundle_path: Path, device: torch.device, p
     # Load model based on file type
     if model_path.suffix == ".ts":
         # TorchScript bundle
-        return torch.jit.load(str(model_path), map_location=device).eval()
+        return cast("torch.nn.Module", torch.jit.load(str(model_path), map_location=device).eval())
     else:
         # .pt checkpoint: instantiate network from config and load state dict
         try:
             # Some .pt files may still be TorchScript; try jit first
-            return torch.jit.load(str(model_path), map_location=device).eval()
+            return cast("torch.nn.Module", torch.jit.load(str(model_path), map_location=device).eval())
         except Exception as ex:
             # Fallback to eager model with loaded weights
             if parser is None:
@@ -128,7 +131,7 @@ def _load_model_from_directory_bundle(bundle_path: Path, device: torch.device, p
                 # Assume raw state dict
                 state_dict = checkpoint
             network.load_state_dict(state_dict, strict=True)
-            return network.eval()
+            return cast("torch.nn.Module", network.eval())
 
 
 def _read_directory_bundle_config(bundle_path_obj: Path, config_names: List[str]) -> ConfigParser: