Update app.py

app.py

@@ -118,8 +118,13 @@ def prediction(df):
     #x_model = tf.keras.models.load_model('my_model.h5')
   try:
         x_model = tf.keras.models.load_model('my_model.h5')
-  except:
-        x_model = tf.keras.layers.TFSMLayer('my_model', call_endpoint='serving_default')
+  except (OSError, ValueError):
+        # Load the model as a TFSMLayer if .h5 loading fails
+        tfs_layer = tf.keras.layers.TFSMLayer('my_model', call_endpoint='serving_default')
+        # Create a new model using the TFSMLayer
+        inputs = tf.keras.Input(shape=(X_test_df.shape[1],))
+        outputs = tfs_layer(inputs)
+        x_model = tf.keras.Model(inputs=inputs, outputs=outputs)
   y_pred = x_model.predict(X_test_df)
   #predicition = []
   #for i in list(y_pred):