code: def zero_grad_ctx(
      parameters: Union[Dict[str, Tensor], Iterable[Tensor]],
      zero_on_enter: bool = True,
      zero_on_exit: bool = False,
  ) -> Generator[None, None, None]:
      def zero_() -> None:
          for param in (
              parameters.values()
              if isinstance(parameters, dict)
              else parameters
          ):
              if param.grad is not None:
                  param.grad.zero_()
  
      if zero_on_enter:
          zero_()
  
      try:
          yield
      finally:
          if zero_on_exit:
              zero_()
  
file_path: botorch/botorch/utils/context_managers.py
imports: from __future__ import annotations
  from contextlib import contextmanager
  from typing import Any, Callable, Dict, Generator, Iterable, NamedTuple, Optional, Union
  from torch import device, dtype, Tensor
  from torch.nn import Module
label_desc: Use image node_id 5 for calling a global function with example usage: zero_grad_ctx(parameters, zero_on_enter, zero_on_exit) without return types