Skip to content

Commit 2234981

Browse files
authored
Fixed pruning examples error (#1127)
1 parent 2651ee4 commit 2234981

File tree

1 file changed

+3
-3
lines changed
  • examples/pytorch/image_recognition/torchvision_models/pruning/magnitude/eager

1 file changed

+3
-3
lines changed

examples/pytorch/image_recognition/torchvision_models/pruning/magnitude/eager/main.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -227,9 +227,9 @@ def train_func(model):
227227
return model
228228

229229
prune.model = common.Model(model)
230-
prune._train_dataloader = train_loader
231-
prune._eval_dataloader = val_loader
232-
prune._pruning_func = train_func
230+
prune.train_dataloader = train_loader
231+
prune.eval_dataloader = val_loader
232+
prune.train_func = train_func
233233
model = prune.fit()
234234
model.save(args.output_model)
235235
return

0 commit comments

Comments
 (0)