Add test

Author: kawakami-masaki0

model_compression_toolkit/core/common/model_collector.py

@@ -256,7 +256,6 @@ def infer(self, inputs_list: List[np.ndarray]):
         hessian_tensors += [None for _ in range(len(activation_tensors) - len(hessian_tensors))]
 
         for activation_tensor, hessian_tensor, stats_container in zip(activation_tensors, hessian_tensors, self.stats_containers_list):
-            print('activation_tensor', activation_tensor.shape)
             if isinstance(stats_container, (list, tuple)):
                 if hessian_tensor is None:
                     hessian_tensor = [None for _ in range(len(activation_tensor))]

tests/pytorch_tests/function_tests/test_torch_utils.py

@@ -29,6 +29,9 @@ def setUp(self):
         self.list_of_numbers = [1, 2, 3]
         self.tuple_of_numbers = (1, 2, 3)
 
+        self.scalar_numpy_array = np.array([1.25])
+        self.scalar_torch_tensor = torch.tensor(1.25)
+
     @patch('model_compression_toolkit.core.pytorch.pytorch_device_config.get_working_device')
     def test_to_torch_tensor_with_numpy_array(self, mock_get_device):
         mock_get_device.return_value = 'cpu'
@@ -69,21 +72,31 @@ def test_torch_tensor_to_numpy_with_torch_tensor(self):
         np.testing.assert_array_almost_equal(result, self.numpy_array)
 
     def test_torch_tensor_to_numpy_with_scalar_tensor(self):
-        scalar_tensor = torch.tensor(1.25)
-        result = torch_tensor_to_numpy(scalar_tensor)
+        result = torch_tensor_to_numpy(self.scalar_torch_tensor)
         self.assertEqual(result.shape, (1,))
-        np.testing.assert_array_almost_equal(result, np.array([1.25]))
+        np.testing.assert_array_almost_equal(result, self.scalar_numpy_array)
 
     def test_torch_tensor_to_numpy_with_list(self):
         result = torch_tensor_to_numpy([self.torch_tensor, self.torch_tensor])
         self.assertEqual(len(result), 2)
         self.assertTrue(all(isinstance(x, np.ndarray) for x in result))
 
+    def test_torch_tensor_to_numpy_with_scalar_list(self):
+        result = torch_tensor_to_numpy([self.scalar_torch_tensor, self.scalar_torch_tensor])
+        self.assertEqual(len(result), 2)
+        self.assertTrue(all(isinstance(x, np.ndarray) for x in result))
+        self.assertTrue(all(x.shape == (1,) for x in result))
+
     def test_torch_tensor_to_numpy_with_tuple(self):
         result = torch_tensor_to_numpy((self.torch_tensor, self.torch_tensor))
         self.assertEqual(len(result), 2)
         self.assertTrue(all(isinstance(x, np.ndarray) for x in result))
 
+    def test_torch_tensor_to_numpy_with_scalar_tuple(self):
+        result = torch_tensor_to_numpy((self.scalar_torch_tensor, self.scalar_torch_tensor))
+        self.assertEqual(len(result), 2)
+        self.assertTrue(all(isinstance(x, np.ndarray) for x in result))
+
     @patch('model_compression_toolkit.logger.Logger')
     def test_torch_tensor_to_numpy_with_unsupported_type(self, mock_logger):
         with self.assertRaises(Exception):