Add method to set model network for nnUNet predictor and update predict method

Author: SimoneBendazzoli93

monai/deploy/operators/monet_bundle_inference_operator.py

@@ -60,12 +60,17 @@ def _init_config(self, config_names):
 
         self._nnunet_predictor = parser.get_parsed_content("network_def")
 
+    def _set_model_network(self, model_network):
+        """Sets the model network for the nnUNet predictor."""
+        if not (isinstance(model_network, torch.nn.Module) or isinstance(model_network, torch.jit.ScriptModule)):
+            raise TypeError(f"Expected model_network to be a torch.nn.Module or a torch.jit.ScriptModule, got {type(model_network)}")
+        self._nnunet_predictor.network = model_network
+
     def predict(self, data: Any, *args, **kwargs) -> Union[Image, Any, Tuple[Any, ...], Dict[Any, Any]]:
         """Predicts output using the inferer. If multimodal data is provided as keyword arguments,
         it concatenates the data with the main input data."""
 
-        self._nnunet_predictor.predictor.network = self._model_network
-
+        self._set_model_network(self._model_network) 
         if len(kwargs) > 0:
             multimodal_data = {"image": data}
             for key in kwargs.keys():