Fix test_compression.py

learning-at-home · Mar 28, 2023 · 67704d7 · 67704d7
1 parent 1920d70
commit 67704d7
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/test_compression.py b/tests/test_compression.py
@@ -49,7 +49,7 @@ def test_tensor_compression(size=(128, 128, 64), alpha=5e-08, beta=0.0008):
 def _check(tensor, compression, rtol=1e-5, atol=1e-8, chunk_size=30 * 1024):
     serialized_tensor = serialize_torch_tensor(tensor, compression)
     chunks = list(split_for_streaming(serialized_tensor, chunk_size))
-    assert len(chunks) == (len(serialized_tensor.buffer) - 1) // chunk_size + 1
+    assert len(chunks) == max((len(serialized_tensor.buffer) - 1) // chunk_size + 1, 1)
     restored = combine_from_streaming(chunks)
     result = deserialize_torch_tensor(restored)
     assert torch.allclose(result, tensor, rtol=rtol, atol=atol)