Fix a small bug in the dependency_aware mode. (#3143)

Signed-off-by: Ningxin <Ningxin.Zheng@microsoft.com>

Fix a small bug in the dependency_aware mode. (#3143)
Signed-off-by: Ningxin <Ningxin.Zheng@microsoft.com>
aea82f23 · Ningxin Zheng · GitHub · 95f731e4 · aea82f23
Unverified Commit aea82f23 authored Dec 02, 2020 by Ningxin Zheng Committed by GitHub Dec 02, 2020
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 1 deletion

nni/algorithms/compression/pytorch/pruning/structured_pruning.py ...orithms/compression/pytorch/pruning/structured_pruning.py +2 -1

No files found.
--- a/nni/algorithms/compression/pytorch/pruning/structured_pruning.py
+++ b/nni/algorithms/compression/pytorch/pruning/structured_pruning.py
@@ -278,7 +278,8 @@ class StructuredWeightMasker(WeightMasker):
                sparsity, _w, _w_idx)
            num_total = current_weight.size(0)
            if num_total < 2 or num_prune < 1:
-                return base_mask
+                masks[name] = base_mask
+                continue
            _tmp_mask = self.get_mask(
                base_mask, current_weight, num_prune, _w, _w_idx, channel_masks)