Update optimization.py

Author: ProGamerGov

captum/optim/_core/optimization.py

@@ -16,7 +16,7 @@
     )
 
 from captum.optim._core.loss import default_loss_summarize
-from captum.optim._core.output_hook import ModuleOutputsHook
+from captum.optim._core.output_hook import ModuleOutputsHook, _remove_all_forward_hooks
 from captum.optim._param.image.images import InputParameterization, NaturalImage
 from captum.optim._param.image.transforms import RandomScale, RandomSpatialJitter
 from captum.optim._utils.typing import (
@@ -196,6 +196,29 @@ def continue_while(
     return continue_while
 
 
+def cleanup_module_hooks(modules: Union[nn.Module, List[nn.Module]) -> None:
+    """
+    Remove any InputOptimization hooks from the specified modules. This may be useful
+    in the event that something goes wrong in between creating the InputOptimization
+    instance and running the optimization function, or if InputOptimization fails
+    without properly removing it's hooks.
+
+    Warning: This function will remove all the hooks placed by InputOptimization
+    instances on the target modules, and thus can interfere with using multiple
+    InputOptimization instances.
+
+    Args:
+
+        modules (nn.Module or list of nn.Module): Any module instances that contain
+            hooks created by InputOptimization, for which the removal of the hooks is
+            required.
+    """
+    if not hasattr(modules, "__iter__"):
+        modules = [modules]
+    # Captum ModuleOutputsHook uses "module_outputs_forward_hook" hook functions
+    [_remove_all_forward_hooks(module, "module_outputs_forward_hook") for module in modules]
+
+
 __all__ = [
     "InputOptimization",
     "n_steps",