code: def test_magnitude_pruner_nondivisible_block_size(
      out_channels, block_size
  ):
      """
      Test block sparsity when the number of channels is not divisible by block size
      """
      conv2d = torch.nn.Conv2d(
          in_channels=3,
          out_channels=out_channels,
          kernel_size=(3, 3),
          bias=False,
          groups=1,
      )
      weight_shape = tuple(conv2d.weight.shape)
      weight_tensor = torch.abs(torch.randn(*weight_shape))
      weight_tensor[weight_tensor == 0] = 1.0
      conv2d.weight = torch.nn.Parameter(weight_tensor)
      config = MagnitudePrunerConfig.from_dict(
          {
              "global_config": {
                  "scheduler": {"update_steps": [1, 2]},
                  "initial_sparsity": 0.0,
                  "target_sparsity": 0.5,
                  "block_size": block_size,
              }
          },
      )
      pruner = MagnitudePruner(conv2d, config)
      conv2d = pruner.prepare()
  
      for _ in range(4):
          pruner.step()
  
      if block_size > 1:
          block_sparse_channels = (
              out_channels - out_channels % block_size
          )
          for idx in range(0, block_sparse_channels, block_size):
              for jdx in range(1, block_size):
                  assert torch.all(
                      conv2d.weight_mask[idx]
                      == conv2d.weight_mask[idx + jdx]
                  )
  
      sparsity = (
          conv2d.weight_mask.eq(0).sum() / conv2d.weight_mask.numel()
      )
      np.testing.assert_array_almost_equal(sparsity, 0.5, decimal=2)
  
imports: import copy
  from collections import OrderedDict
  import numpy
  import pytest
  import torch
  from coremltools.optimize.torch.pruning import MagnitudePruner, MagnitudePrunerConfig, ModuleMagnitudePrunerConfig
  from coremltools.optimize.torch.pruning._utils import n_m_mask
label_desc: Use image node_id 8 for calling a global function with example usage: test_magnitude_pruner_nondivisible_block_size(out_channels, block_size) without return types