AttributeError: partially initialized module 'torch_pruning' has no attribute 'DependencyGraph'

[nbansal90]

I am just beginning to setup the repository. Getting this error running the sample code provided in the Readme.

AttributeError: partially initialized module 'torch_pruning' has no attribute 'DependencyGraph'

CODE RUNNING

import torch
from torchvision.models import resnet18
import torch_pruning as tp

model = resnet18(pretrained=True).eval()

# 1. Build dependency graph for resnet18
DG = tp.DependencyGraph().build_dependency(model, example_inputs=torch.randn(1,3,224,224))

# 2. Group coupled layers for model.conv1
group = DG.get_pruning_group( model.conv1, tp.prune_conv_out_channels, idxs=[2, 6, 9] )

# 3. Prune grouped layers altogether
if DG.check_pruning_group(group): # avoid full pruning, i.e., channels=0.
    group.prune()

# 4. Save & Load
model.zero_grad() # clear gradients
torch.save(model, 'model.pth') # We can not use .state_dict as the model structure is changed.
model = torch.load('model.pth') # load the pruned model

Any help/suggestions greatly appreciated!

[janthmueller]

Do you have another file/module called torch_pruning in your project? This might be the problem.