Creating a Pruning Runner - 2.5 English

Import modules and prepare input signature:

from pytorch_nndct import get_pruning_runner

# The input signature should have the same shape and dtype as the model input.
input_signature = torch.randn([1, 3, 224, 224], dtype=torch.float32) 

Create an iterative pruning runner:

runner = get_pruning_runner(model, input_signature, 'iterative')

Or, a one-step pruning runner:

runner = get_pruning_runner(model, input_signature, 'one_step')