fix keras extension in integ test (#4834)

pintaoz-aws · web-flow · commit 686d6f05c75d · 2024-08-14T15:24:19.000-07:00
diff --git a/tests/data/tensorflow_mnist/mnist_v2.py b/tests/data/tensorflow_mnist/mnist_v2.py
@@ -198,7 +198,7 @@ def main(args):
 
         if args.current_host == args.hosts[0]:
             ckpt_manager.save()
-            net.save("/opt/ml/model/1")
+            net.save("/opt/ml/model/1.keras")
 
 
 if __name__ == "__main__":
