Change to net.export() for TF2.16

Author: shantanutrip

tests/data/tensorflow_mnist/mnist_v2.py

@@ -198,7 +198,10 @@ def main(args):
 
         if args.current_host == args.hosts[0]:
             ckpt_manager.save()
-            net.save("/opt/ml/model/1")
+            if int(tf_major) > 2 or (int(tf_major) == 2 and int(tf_minor) >= 16):
+                net.export("/opt/ml/model/1")
+            else:
+                net.save("/opt/ml/model/1")
 
 
 if __name__ == "__main__":