fix jax failing tests (#2231)

keras-team · Dec 11, 2023 · 1f6bf4a · 1f6bf4a
1 parent 467f19d
commit 1f6bf4a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/keras_cv/layers/regularization/drop_path.py b/keras_cv/layers/regularization/drop_path.py
@@ -50,7 +50,7 @@ class DropPath(keras.layers.Layer):
     def __init__(self, rate=0.5, seed=None, **kwargs):
         super().__init__(**kwargs)
         self.rate = rate
-        self.seed = seed
+        self.seed = keras.random.SeedGenerator(seed=seed)
 
     def call(self, x, training=None):
         if self.rate == 0.0 or not training: