def test_pytorch_model_save():
artifacts_folder = os.path.join(TESTS_PATH, "serve", "data",
"mlflow_pytorch", "model")
pytorch_mnist_model = Model(
name="test-pytorch-mnist",
platform=ModelFramework.MLFlow,
local_folder=artifacts_folder,
description="A pytorch MNIST model - python 3.7",
)
pytorch_mnist_model.save(save_env=True)
remote_model = deploy_local(pytorch_mnist_model)
data = np.random.randn(1, 28 * 28).astype(np.float32)
remote_model.predict(data)
remote_model.undeploy()
def test_model_save(custom_model: Model):
custom_model.save(save_env=False)
loaded = Model.load(custom_model.details.local_folder)
assert len(custom_model.context.__dict__) > 0
assert len(loaded.context.__dict__) == 0