hotfix for pytorch missing locations

nncf/torch/sparsity/const/algo.py
nncf/torch/sparsity/magnitude/algo.py
nncf/torch/sparsity/rb/algo.py

Author: shumaari

nncf/torch/sparsity/const/algo.py

@@ -11,6 +11,7 @@
 from typing import Tuple
 
 from nncf.common.graph import NNCFNode
+from nncf.common.graph.utils import get_weight_shape_legacy
 from nncf.common.sparsity.statistics import ConstSparsityStatistics
 from nncf.common.statistics import NNCFStatistics
 from nncf.common.utils.api_marker import api
@@ -26,7 +27,7 @@
 @PT_COMPRESSION_ALGORITHMS.register("const_sparsity")
 class ConstSparsityBuilder(BaseSparsityAlgoBuilder):
     def create_weight_sparsifying_operation(self, target_module_node: NNCFNode, compression_lr_multiplier: float):
-        return BinaryMask(target_module_node.layer_attributes.get_weight_shape())
+        return BinaryMask(get_weight_shape_legacy(target_module_node.layer_attributes))
 
     def _build_controller(self, model: NNCFNetwork) -> PTCompressionAlgorithmController:
         return ConstSparsityController(model, self._sparsified_module_info)

nncf/torch/sparsity/magnitude/algo.py

@@ -17,6 +17,7 @@
 from nncf.api.compression import CompressionStage
 from nncf.common.accuracy_aware_training.training_loop import ADAPTIVE_COMPRESSION_CONTROLLERS
 from nncf.common.graph import NNCFNode
+from nncf.common.graph.utils import get_weight_shape_legacy
 from nncf.common.initialization.batchnorm_adaptation import BatchnormAdaptationAlgorithm
 from nncf.common.schedulers import StubCompressionScheduler
 from nncf.common.sparsity.schedulers import SPARSITY_SCHEDULERS
@@ -44,7 +45,7 @@
 @PT_COMPRESSION_ALGORITHMS.register("magnitude_sparsity")
 class MagnitudeSparsityBuilder(BaseSparsityAlgoBuilder):
     def create_weight_sparsifying_operation(self, target_module_node: NNCFNode, compression_lr_multiplier: float):
-        return BinaryMask(target_module_node.layer_attributes.get_weight_shape())
+        return BinaryMask(get_weight_shape_legacy(target_module_node.layer_attributes))
 
     def _build_controller(self, model: NNCFNetwork) -> PTCompressionAlgorithmController:
         return MagnitudeSparsityController(model, self._sparsified_module_info, self.config)

nncf/torch/sparsity/rb/algo.py

@@ -18,6 +18,7 @@
 from nncf.api.compression import CompressionStage
 from nncf.common.accuracy_aware_training.training_loop import ADAPTIVE_COMPRESSION_CONTROLLERS
 from nncf.common.graph import NNCFNode
+from nncf.common.graph.utils import get_weight_shape_legacy
 from nncf.common.schedulers import StubCompressionScheduler
 from nncf.common.sparsity.schedulers import SPARSITY_SCHEDULERS
 from nncf.common.sparsity.statistics import RBSparsityStatistics
@@ -44,7 +45,7 @@
 class RBSparsityBuilder(BaseSparsityAlgoBuilder):
     def create_weight_sparsifying_operation(self, target_module_node: NNCFNode, compression_lr_multiplier: float):
         return RBSparsifyingWeight(
-            target_module_node.layer_attributes.get_weight_shape(),
+            get_weight_shape_legacy(target_module_node.layer_attributes),
             frozen=False,
             compression_lr_multiplier=compression_lr_multiplier,
         )