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)

    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.
    
coremltools/test/optimize/torch/pruning/test_magnitude_pruner.py