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.