Update app.py

app.py

@@ -8,19 +8,23 @@ from tensorflow.keras import backend as K
 
 # Define the custom FixedDropout layer
 class FixedDropout(tf.keras.layers.Layer):
-    def __init__(self, rate, **kwargs):
+    def __init__(self, rate, noise_shape=None, **kwargs):
         super(FixedDropout, self).__init__(**kwargs)
         self.rate = rate
+        self.noise_shape = noise_shape  # Include the noise_shape argument
 
     def call(self, inputs, training=None):
         if training is None:
             training = K.learning_phase()
-        return K.in_train_phase(K.dropout(inputs, self.rate), inputs, training=training)
+        return K.in_train_phase(K.dropout(inputs, self.rate, noise_shape=self.noise_shape), inputs, training=training)
 
     def get_config(self):
         config = super(FixedDropout, self).get_config()
+        config['rate'] = self.rate  # Serialize the rate argument
+        config['noise_shape'] = self.noise_shape  # Serialize the noise_shape argument
         return config
 
+
 class ImageClassifierApp:
     def __init__(self, model_path):
         self.model_path = model_path