Update train.py

train.py

@@ -86,13 +86,16 @@ print('Test accuracy:', test_acc)
 # Save the model for future use (optional)
 # Not recommended for Hugging Face Hub upload (use tf.saved_model.save())
 
-# # Export the model for Hugging Face Hub using tf.saved_model.save()
-# export_dir = 'saved_model'  # Create a directory for the SavedModel
-# tf.saved_model.save(model, export_dir)
+# Export the model for Hugging Face Hub using tf.saved_model.save()
+export_dir = 'saved_model'  # Create a directory for the SavedModel
+tf.saved_model.save(model, export_dir)
+Save model weights to local directory
+# model.save_pretrained("kerascatanddog")
 
 # Upload the model to your Hugging Face space repository
 push_to_hub_keras(
-    model,  # Point to the SavedModel directory
+    model=export_dir,
+    # model,  # Point to the SavedModel directory
     repo_id="okeowo1014/kerascatanddog",
     commit_message="cats and dog image classifier with transfer learning",
     tags=["image-classifier", "data-augmentation", "class-weights"],