code: def test_fabric_optimizer_zero_grad_kwargs():
      """Test that Fabric can adapt the `.zero_grad()` arguments to the underlying optimizer."""
      # Test PyTorch's standard `.zero_grad()` signature
      with mock.patch("torch.optim.SGD.zero_grad") as zero_grad_mock:
          optimizer = torch.optim.SGD(
              torch.nn.Linear(1, 1).parameters(), 0.1
          )
          fabric_optimizer = _FabricOptimizer(
              optimizer=optimizer, strategy=Mock()
          )
          fabric_optimizer.zero_grad()
          zero_grad_mock.assert_called_with()
          fabric_optimizer.zero_grad(set_to_none=False)
          zero_grad_mock.assert_called_with(set_to_none=False)
          fabric_optimizer.zero_grad(set_to_none=True)
          zero_grad_mock.assert_called_with(set_to_none=True)
  
      # Test weird `.zero_grad()` signatures from other libraries
      custom_zero_grad = Mock()
  
      class CustomSGD(torch.optim.SGD):
          def zero_grad(self, set_grads_to_None=False):
              custom_zero_grad(set_grads_to_None=set_grads_to_None)
  
      optimizer = CustomSGD(torch.nn.Linear(1, 1).parameters(), 0.1)
      fabric_optimizer = _FabricOptimizer(
          optimizer=optimizer, strategy=Mock()
      )
      fabric_optimizer.zero_grad()
      custom_zero_grad.assert_called_with(set_grads_to_None=False)
  
imports: from unittest import mock
  from unittest.mock import Mock, call
  import pytest
  import torch
  from lightning.fabric.fabric import Fabric
  from lightning.fabric.plugins import Precision
  from lightning.fabric.utilities.device_dtype_mixin import _DeviceDtypeModuleMixin
  from lightning.fabric.utilities.imports import _TORCH_GREATER_EQUAL_2_1
  from lightning.fabric.wrappers import _FabricDataLoader, _FabricModule, _FabricOptimizer, _unwrap_objects, is_wrapped
  from torch.utils.data import BatchSampler, DistributedSampler
  from torch.utils.data.dataloader import DataLoader
  from tests_fabric.helpers.runif import RunIf
label_desc: Use image node_id 15 for calling a global function with example usage: test_fabric_optimizer_zero_grad_kwargs() without return types