package unsupported ml model

predict.py

@@ -7,7 +7,7 @@ import torch
 import mlflow
 
 
-if __name__ == '__main__':
+def main():
 
     # set MLflow server
     mlflow.set_tracking_uri('http://127.0.0.1:5001')
@@ -21,4 +21,8 @@ if __name__ == '__main__':
     print(my_fortune)
 
     # save model and env to local file system
-    mlflow.pytorch.save_model(model, './fortune_predict_model')
+    mlflow.pytorch.save_model(model, './model')
+
+
+if __name__ == '__main__':
+    main()