code: def test_no_grad_context():
      trainer = Mock()
  
      class Foo:
          @_no_grad_context
          def run(self):
              ...
  
      f = Foo()
      with pytest.raises(TypeError, match="Foo` needs to be a Loop"):
          f.run()
  
      class Foo(_Loop):
          @_no_grad_context
          def run(self):
              ...
  
      f = Foo(trainer)
      with pytest.raises(
          TypeError, match="Foo.inference_mode` needs to be defined"
      ):
          f.run()
  
      class Foo(_Loop):
          def __init__(self):
              super().__init__(trainer)
              self.inference_mode = False
  
          @_no_grad_context
          def run(self):
              ...
  
      f = Foo()
      with mock.patch("torch.no_grad") as no_grad_mock:
          f.run()
      no_grad_mock.assert_called_once_with()
      f.inference_mode = True
      with mock.patch("torch.inference_mode") as inference_mode_mock:
          f.run()
      if not _TORCH_EQUAL_2_0:
          inference_mode_mock.assert_called_once_with()
  
file_path: pytorch-lightning/tests/tests_pytorch/trainer/flags/test_inference_mode.py
imports: from unittest import mock
  from unittest.mock import Mock
  import pytest
  import torch
  from lightning.fabric.utilities.imports import _TORCH_EQUAL_2_0
  from lightning.pytorch import Trainer
  from lightning.pytorch.demos.boring_classes import BoringModel
  from lightning.pytorch.loops import _Loop
  from lightning.pytorch.loops.utilities import _no_grad_context
label_desc: Use image node_id 2 for calling a global function with example usage: test_no_grad_context() without return types