The pruner requires two arguments:
- The model to be pruned
- The inputs needed by the model for inference
import torch from pytorch_nndct import OFAPruner inputs = torch.randn([1, 3, 224, 224], dtype=torch.float32) pruner = OFAPruner(model, inputs)
Note: The input does not need to be real data. You can use randomly generated dummy data if it has the same shape and type as the real data.