A pruner can be created by providing the model to be pruned and its input
shape and input dtype. Note that shape is the size of the input image and does not
contain batch size.
from pytorch_nndct import Pruner
from pytorch_nndct import InputSpec
pruner = Pruner(model, InputSpec(shape=(3, 224, 224), dtype=torch.float32))
For models with multiple inputs, you can use a list of InputSpec
to initialize a pruner.