code review modification 2

sony · Oct 30, 2023 · a85cb79 · a85cb79
1 parent 1b88e58
commit a85cb79
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/model_compression_toolkit/core/pytorch/hessian/weights_trace_hessian_calculator_pytorch.py b/model_compression_toolkit/core/pytorch/hessian/weights_trace_hessian_calculator_pytorch.py
@@ -124,5 +124,9 @@ def compute(self) -> np.ndarray:
         # Compute the mean of the approximations
         final_approx = torch.mean(torch.stack(approximation_per_iteration), dim=0)
 
+        # Make sure all final shape are tensors and not scalar
+        if self.hessian_request.granularity == HessianInfoGranularity.PER_TENSOR:
+            final_approx = final_approx.reshape(1)
+
         return final_approx.detach().cpu().numpy()
 
diff --git a/tests/pytorch_tests/function_tests/test_hessian_info_weights.py b/tests/pytorch_tests/function_tests/test_hessian_info_weights.py
@@ -122,7 +122,7 @@ def get_expected_shape(weights_shape, granularity):
     if granularity==hessian_common.HessianInfoGranularity.PER_ELEMENT:
         return weights_shape
     elif granularity==hessian_common.HessianInfoGranularity.PER_TENSOR:
-        return ()
+        return (1,)
     else:
         return (weights_shape[0],)