Allow gaussian_blur sigma to be a jax.Array.

PiperOrigin-RevId: 645018560
google-deepmind · Jun 20, 2024 · 2b741ee · 2b741ee
1 parent 07e9cb7
commit 2b741ee
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/dm_pix/_src/augment.py b/dm_pix/_src/augment.py
@@ -466,7 +466,7 @@ def flip_up_down(
 
 def gaussian_blur(
     image: chex.Array,
-    sigma: float,
+    sigma: chex.Numeric,
     kernel_size: float,
     *,
     padding: str = "SAME",