fix slice.

Author: sayakpaul

tests/quantization/bnb/test_4bit.py

@@ -29,7 +29,6 @@
     is_transformers_available,
     load_pt,
     numpy_cosine_similarity_distance,
-    print_tensor_test,
     require_accelerate,
     require_bitsandbytes_version_greater,
     require_torch,
@@ -433,7 +432,6 @@ def test_quality(self):
         expected_slice = np.array([0.1123, 0.1296, 0.1609, 0.1042, 0.1230, 0.1274, 0.0928, 0.1165, 0.1216])
 
         max_diff = numpy_cosine_similarity_distance(expected_slice, out_slice)
-        print_tensor_test(out_slice)
         self.assertTrue(max_diff < 1e-2)
 
     def test_generate_quality_dequantize(self):
@@ -449,7 +447,6 @@ def test_generate_quality_dequantize(self):
         ).images
 
         out_slice = output[0, -3:, -3:, -1].flatten()
-        print_tensor_test(out_slice)
         expected_slice = np.array([0.1216, 0.1387, 0.1584, 0.1152, 0.1318, 0.1282, 0.1062, 0.1226, 0.1228])
         max_diff = numpy_cosine_similarity_distance(expected_slice, out_slice)
         self.assertTrue(max_diff < 1e-3)
@@ -523,7 +520,6 @@ def test_quality(self):
         ).images
 
         out_slice = output[0, -3:, -3:, -1].flatten()
-        print_tensor_test(out_slice)
         expected_slice = np.array([0.0583, 0.0586, 0.0632, 0.0815, 0.0813, 0.0947, 0.1040, 0.1145, 0.1265])
 
         max_diff = numpy_cosine_similarity_distance(expected_slice, out_slice)

tests/quantization/bnb/test_mixed_int8.py

@@ -26,7 +26,6 @@
     is_transformers_available,
     load_pt,
     numpy_cosine_similarity_distance,
-    print_tensor_test,
     require_accelerate,
     require_bitsandbytes_version_greater,
     require_torch,
@@ -370,8 +369,7 @@ def test_quality(self):
             output_type="np",
         ).images
         out_slice = output[0, -3:, -3:, -1].flatten()
-        print_tensor_test(out_slice)
-        expected_slice = np.array([0.0149, 0.0322, 0.0073, 0.0134, 0.0332, 0.011, 0.002, 0.0232, 0.0193])
+        expected_slice = np.array([0.0376, 0.0359, 0.0015, 0.0449, 0.0479, 0.0098, 0.0083, 0.0295, 0.0295])
 
         max_diff = numpy_cosine_similarity_distance(expected_slice, out_slice)
         self.assertTrue(max_diff < 1e-2)
@@ -420,7 +418,6 @@ def test_generate_quality_dequantize(self):
         ).images
 
         out_slice = output[0, -3:, -3:, -1].flatten()
-        print_tensor_test(out_slice)
         expected_slice = np.array([0.0266, 0.0264, 0.0271, 0.0110, 0.0310, 0.0098, 0.0078, 0.0256, 0.0208])
         max_diff = numpy_cosine_similarity_distance(expected_slice, out_slice)
         self.assertTrue(max_diff < 1e-2)
@@ -471,7 +468,6 @@ def test_quality(self):
             output_type="np",
         ).images
         out_slice = output[0, -3:, -3:, -1].flatten()
-        print_tensor_test(out_slice)
         expected_slice = np.array([0.0574, 0.0554, 0.0581, 0.0686, 0.0676, 0.0759, 0.0757, 0.0803, 0.0930])
 
         max_diff = numpy_cosine_similarity_distance(expected_slice, out_slice)