test: ensure float32 works as a dtype input for sharded

seung-lab · Jan 29, 2025 · 614a1fb · 614a1fb
1 parent e5025d1
commit 614a1fb
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/test/test_shards.py b/test/test_shards.py
@@ -44,16 +44,17 @@ def prod(x):
   },
 ]
 
+@pytest.mark.parametrize("dtype", [np.uint8, np.float32])
 @pytest.mark.parametrize("scale", SCALES)
-def test_sharded_image_bits(scale):
+def test_sharded_image_bits(scale, dtype):
   dataset_size = Vec(*scale["size"])
   chunk_size = Vec(*scale["chunk_sizes"][0])
 
   spec = create_sharded_image_info( 
     dataset_size=dataset_size,
     chunk_size=chunk_size,
     encoding=scale["encoding"],
-    dtype=np.uint8
+    dtype=dtype
   )
 
   shape = image_shard_shape_from_spec(