code: def test_magnitude_pruning_training_and_validation(
      simple_module, sample_data
  ):
      """
      Tests pruned weights are used for computing forward pass
      pruned module. Also demonstrates how pruner can be combined with
      training code in PyTorch, i.e, pruning can be done at a schedule
      different from training.
  
      Note: No actual training happens here because loss function is a no-op.
      """
      config = MagnitudePrunerConfig.from_dict(
          {
              "global_config": {
                  "scheduler": {"update_steps": [2, 3]},
                  "initial_sparsity": 0.0,
                  "target_sparsity": 0.75,
                  "granularity": "per_channel",
              }
          },
      )
      pruner = MagnitudePruner(simple_module, config)
      simple_module = pruner.prepare()
  
      # Train the model for 4 epochs
      num_epochs = 4
      X, Y = sample_data
      simple_module.train()
      optimizer = torch.optim.Adam(
          params=simple_module.parameters(), lr=0.001
      )
      for _ in range(num_epochs):
          for inp, label in zip(X, Y):
              inp = inp.view(1, *X.shape[1:])
              label = label.view(1, *Y.shape[1:])
              output = simple_module(inp)
              loss = _zero_loss(output, label)
              optimizer.zero_grad()
              loss.backward()
              optimizer.step()
          pruner.step()
  
      # Test inference
      # After 4 iterations, pruner will zero out first 3 layers of conv2d layer in simple_module
      simple_module.eval()
      with torch.no_grad():
          x_test = torch.tensor(
              np.random.uniform(0.0, 1.0, size=(1, 3, 8, 8)).astype(
                  np.float32
              )
          )
          y_test = simple_module(x_test).detach().numpy()
  
      zero_output = y_test[:, :3, :, :]
      nonzero_output = y_test[:, 3:, :, :]
      np.testing.assert_equal(zero_output, np.zeros_like(zero_output))
      assert np.any(np.abs(nonzero_output) > 0.0)
  
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 15 for calling a global function with example usage: test_magnitude_pruning_training_and_validation(simple_module, sample_data) without return types