Repository: intel/ipex-llm
Pull request: [Nano] consistent save/load for nano pytorch
resolves #4485 

Problem statement:
[Nano] trainer.save accepts multiple types of models
# Description
Current API for model saving and loading. It can do:
1. save/export a pytorch model as openvino format 
2. save a `PytorchOpenVINOModel ` as local openvino file

**Do you think this is a bit confusing for users accepting multiple types of models?**   @jason-dai   
Otherwise we can just keep one of them.
### `trainer.save`
save from `PytorchOpenvinoModel`
```python
openvino_model = trainer.trace(model: nn.Module, x, accelerator='openvino')
trainer.save(openvino_model, 'model.xml`)
```
or save from torch model directly:
```python
trainer.save(model: nn.Module, 'model.xml`, input_sample=x, accelerator='openvino', input_sample=x)
```

### `trainer.load`
```python
openvino_model  = PytorchOpenVINOModel.load('model.xml')
```
```python
openvino_model = Trainer.load('saved_openvino_model.xml', accelerator='openvino')
```

Hints from issue discussion:
We should have only one `save` method, which exports standard model format (e.g., OpenVINO or ONNX format) that can be loaded by both nano and standard 3rd party tools (such as OpenVINO or ONNX Runtime). 
I think for `trainer.save`, we may follow your first design (**and only support this design**), which means that our users will call `trace` or `quantize` to an accelerated model and then call `save` to dump a/multiple corresponding files(`.onnx`/`.xml`/`.state`...)

For pytorch model, the `trainer.load` could be tricky. Users could not take the dump file to a totally new runtime environment and load it(theoretically impossible). They will need a `model` instance as input. So I think we do need this model parameter. https://github.com/intel-analytics/BigDL/pull/4418#discussion_r857090273
> I think for `trainer.save`, we may follow your first design (**and only support this design**), which means that our users will call `trace` or `quantize` to an accelerated model and then call `save` to dump a/multiple corresponding files(`.onnx`/`.xml`/`.state`...)
> 
> For pytorch model, the `trainer.load` could be tricky. Users could not take the dump file to a totally new runtime environment and load it(theoretically impossible). They will need a `model` instance as input. So I think we do need this model parameter. [#4418 (comment)](https://github.com/intel-analytics/BigDL/pull/4418#discussion_r857090273)

What should the workflow look like to save/load a PyTorch model using the Nano API then?
### For OpenVINO/ ONNX
save
```python
qmodel = trainer.quantize(model, accelerator='openvino/onnx')
trainer.save(qmodel, path)
```
load
```python
loaded_qmodel = trainer.load(path)
```
### For Pytorch
save when training finished
```python
model = resnet(..)
qmodel = trainer.quantize(model, accelerator=None)
trainer.save(qmodel, path)
```
load to run on other machine
```python
model = resnet(..)
loaded_qmodel = trainer.load(path, model)
```

LGTM

Commit message: remove duplicate import

Patch to review:
diff --git a/python/nano/src/bigdl/nano/deps/openvino/core/openvino_model.py b/python/nano/src/bigdl/nano/deps/openvino/core/openvino_model.py
index b6ca0440f89..19022268b79 100644
--- a/python/nano/src/bigdl/nano/deps/openvino/core/openvino_model.py
+++ b/python/nano/src/bigdl/nano/deps/openvino/core/openvino_model.py
@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
+from pathlib import Path
 from openvino.inference_engine import IECore
 
 
@@ -39,12 +40,13 @@ def compile_executable(self, inputs):
     def read_network(self, model: str):
         self.ie_network = IECore().read_network(model=model)
 
-    def save(self, path):
+    def _save_model(self, path):
         """
         Save PytorchOpenVINOModel to local as xml and bin file
 
         :param path: Path to save the model.
         """
+        path = Path(path)
         assert self.ie_network, "self.ie_network shouldn't be None."
-        assert path.split('.')[-1] == "xml", "Path of openvino model must be with '.xml' suffix."
-        self.ie_network.serialize(path)
+        assert path.suffix == ".xml", "Path of openvino model must be with '.xml' suffix."
+        self.ie_network.serialize(str(path))
diff --git a/python/nano/src/bigdl/nano/deps/openvino/openvino_api.py b/python/nano/src/bigdl/nano/deps/openvino/openvino_api.py
index cbcd66a3340..7ff754b627b 100644
--- a/python/nano/src/bigdl/nano/deps/openvino/openvino_api.py
+++ b/python/nano/src/bigdl/nano/deps/openvino/openvino_api.py
@@ -28,3 +28,8 @@ def PytorchOpenVINOModel(model, input_sample=None):
     """
     from .pytorch.pytorch_openvino_model import PytorchOpenVINOModel
     return PytorchOpenVINOModel(model, input_sample)
+
+
+def load_openvino_model(path):
+    from .pytorch.pytorch_openvino_model import PytorchOpenVINOModel
+    return PytorchOpenVINOModel._load(path)
diff --git a/python/nano/src/bigdl/nano/deps/openvino/pytorch/pytorch_openvino_model.py b/python/nano/src/bigdl/nano/deps/openvino/pytorch/pytorch_openvino_model.py
index d0c5cdb573f..9d07527238e 100644
--- a/python/nano/src/bigdl/nano/deps/openvino/pytorch/pytorch_openvino_model.py
+++ b/python/nano/src/bigdl/nano/deps/openvino/pytorch/pytorch_openvino_model.py
@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-import os
+from pathlib import Path
 from ..core.openvino_model import OpenVINOModel
 from bigdl.nano.utils.inference.pytorch.model import AcceleratedLightningModule
 from .pytorch_openvino_utils import export
@@ -35,10 +35,11 @@ def __init__(self, model, input_sample=None):
         if isinstance(model, torch.nn.Module):
             export(model, input_sample, 'tmp.xml')
             ov_model_path = 'tmp.xml'
-        AcceleratedLightningModule.__init__(self, None)
         OpenVINOModel.__init__(self, ov_model_path)
-        if os.path.exists('tmp.xml'):
-            os.remove('tmp.xml')
+        AcceleratedLightningModule.__init__(self, None)
+        xml_path = Path('tmp.xml')
+        if xml_path.exists():
+            xml_path.unlink()
 
     def on_forward_start(self, inputs):
         if self.ie_network is None:
@@ -52,13 +53,29 @@ def on_forward_end(self, outputs):
         outputs = self.numpy_to_tensors(outputs.values())
         return outputs
 
+    @property
+    def status(self):
+        status = super().status
+        status.update({"xml_path": 'ov_saved_model.xml', "weight_path": 'ov_saved_model.bin'})
+        return status
+
     @staticmethod
-    def load(path):
+    def _load(path):
         """
-        Load an OpenVINO model for inference.
+        Load an OpenVINO model for inference from directory.
 
         :param path: Path to model to be loaded.
         :return: PytorchOpenVINOModel model for OpenVINO inference.
         """
-        assert path.split('.')[-1] == "xml", "Path of openvino model must be with '.xml' suffix."
-        return PytorchOpenVINOModel(path)
+        status = PytorchOpenVINOModel._load_status(path)
+        if status.get('xml_path', None):
+            xml_path = Path(status['xml_path'])
+            assert xml_path.suffix == '.xml', "Path of openvino model must be with '.xml' suffix."
+        else:
+            raise KeyError("meta-data.yml must specify 'xml_path' for loading.")
+        xml_path = Path(path) / status['xml_path']
+        return PytorchOpenVINOModel(xml_path)
+
+    def _save_model(self, path):
+        xml_path = Path(path) / self.status['xml_path']
+        super()._save_model(xml_path)
diff --git a/python/nano/src/bigdl/nano/pytorch/trainer/Trainer.py b/python/nano/src/bigdl/nano/pytorch/trainer/Trainer.py
index 5e3ed5b577c..2b60db76048 100644
--- a/python/nano/src/bigdl/nano/pytorch/trainer/Trainer.py
+++ b/python/nano/src/bigdl/nano/pytorch/trainer/Trainer.py
@@ -13,11 +13,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
+import copy
 from logging import warning
-from operator import xor
-import os
+from pathlib import Path
 from typing import Any, List, Optional
-
 import pytorch_lightning as pl
 import torch
 from pytorch_lightning import LightningModule
@@ -28,12 +27,13 @@
 from torch.utils.data import DataLoader
 from torchmetrics.metric import Metric
 from torch.optim.lr_scheduler import _LRScheduler
+import yaml
 from bigdl.nano.common import check_avx512
 from bigdl.nano.pytorch.lightning import LightningModuleFromTorch
 from bigdl.nano.pytorch.plugins.ddp_spawn import DDPSpawnPlugin
 from bigdl.nano.deps.ray.ray_api import distributed_ray
 from bigdl.nano.deps.ipex.ipex_api import create_IPEXAccelerator, ipex_device
-from bigdl.nano.deps.openvino.openvino_api import PytorchOpenVINOModel
+from bigdl.nano.deps.openvino.openvino_api import PytorchOpenVINOModel, load_openvino_model
 from bigdl.nano.deps.onnxruntime.onnxruntime_api import bind_onnxrt_methods
 
 distributed_backends = ["spawn", "ray", "subprocess"]
@@ -310,56 +310,74 @@ def trace(model: nn.Module, input_sample=None, accelerator=None):
                              model is a LightningModule with any dataloader attached.
         :param accelerator: The accelerator to use, defaults to None meaning staying in Pytorch
                             backend. Only 'openvino' is supported for now.
-        :return: A PytorchOpenVINOModel using openvino for inference if accelerator='openvino'
+        :return: Model with different acceleration(OpenVINO/ONNX).
         """
         if accelerator == 'openvino':
             return PytorchOpenVINOModel(model, input_sample)
 
     @staticmethod
-    def save(model, path, precision=None, accelerator=None, input_sample=None):
+    def save(model: LightningModule, path):
         """
-        Save the model to path with desired precision and format(by assigning 'accelerator').
-
-        Using a pytorch model, you can export to ONNX, OpenVINO directly with
-        accelerator='openvino'/'onnx'.
+        Save the model to local file.
 
-        :param model: Any model of torch.nn.Module, including PytorchOpenVINOModel
-        :param path: Path to saved model. You need to specify path suffix carefully.
-                     For example, 'model.xml' for OpenVINO, 'model.onnx' for ONNX.
-        :param precision: Precision of saved model, 'FP32', 'FP16', 'INT8', defaults to None,
-                          then the precision depends on the precision of model.
-        :param accelerator: Saved model format and its future accelerator, defaults to None.
-                            accelerator=None will save Pytorch model, accelerator='openvino' will
-                            save an IR to inference by OpenVINO. Same as ONNX.
-        :param input_sample: A set of inputs for trace, defaults to None if you have trace before or
-                             model is a LightningModule with any dataloader attached,
-                             defaults to None.
-                             Only ONNX and OpenVINO requires an input sample to export the model.
+        :param model: Any model of torch.nn.Module, including PytorchOpenVINOModel,
+         PytorchONNXModel.
+        :param path: Path to saved model. Path should be a directory.
         """
-        if accelerator == 'openvino':
-            if precision is None:
-                if not hasattr(model, "save") and isinstance(model, nn.Module):
-                    model = PytorchOpenVINOModel(model, input_sample=input_sample)
-                else:
-                    raise TypeError("Model type of {} can not be exported.".format(type(model)))
-                model.save(path)
+        path = Path(path)
+        Path.mkdir(path, exist_ok=True)
+        if hasattr(model, '_save'):
+            model._save(path)
+        else:
+            # typically for models of nn.Module, LightningModule and LightningModuleFromTorch type
+            meta_path = Path(path) / "meta-data.yml"
+            with open(meta_path, 'w+') as f:
+                metadata = {
+                    'ModelType': 'PytorchModel',
+                    'checkpoint': 'saved_weight.pt'
+                }
+                yaml.safe_dump(metadata, f)
+            checkpoint_path = path / metadata['checkpoint']
+            torch.save(model.state_dict(), checkpoint_path)
 
     @staticmethod
-    def load(path, accelerator=None):
+    def load(path, model: LightningModule = None):
         """
         Load a model from local.
 
-        :param path: Path to model to be loaded. You need to specify path suffix carefully.
-                     For example, 'model.xml' for OpenVINO, 'model.onnx' for ONNX.
-        :param accelerator: Saved model format, defaults to None.
-                            accelerator=None will save Pytorch model, accelerator='openvino' will
-                            save an IR to inference by OpenVINO. Same as ONNX.
-        :return: A PytorchOpenVINOModel using openvino for inference if accelerator='openvino'
+        :param path: Path to model to be loaded. Path should be a directory.
+        :param model: Required FP32 model to load pytorch model. Optional for ONNX/OpenVINO.
+        :return: Model with different acceleration(None/OpenVINO/ONNX) or
+                 precision(FP32/FP16/BF16/INT8).
         """
-        if accelerator == 'openvino' or path.split('.')[-1] == 'xml':
-            # TODO: Need to fix this with lazy import class.
-            # The usage should be:
-            #   PytorchOpenVINOModel.load(path)
-            if not os.path.exists(path):
-                raise FileNotFoundError("{} doesn't exist.".format(path))
-            return PytorchOpenVINOModel(path)
+        path = Path(path)
+        if not path.exists():
+            raise FileNotFoundError("{} doesn't exist.".format(path))
+        meta_path = path / "meta-data.yml"
+        if not meta_path.exists():
+            raise FileNotFoundError("File {} is required to load model.".format(str(meta_path)))
+        with open(meta_path, 'r') as f:
+            metadata = yaml.safe_load(f)
+        model_type = metadata.get('ModelType', None)
+        if model_type == 'PytorchOpenVINOModel':
+            assert model is None, "Argument 'model' must be None for OpenVINO loading."
+            return load_openvino_model(path)
+        # if model_type == 'PytorchONNXModel':
+        # if model_type == 'PytorchQuantizedModel':
+        # ... to be implemented
+        if isinstance(model, nn.Module):
+            # typically for models of nn.Module, LightningModule and LightningModuleFromTorch type
+            model = copy.deepcopy(model)
+            checkpoint_path = metadata.get('checkpoint', None)
+            if checkpoint_path:
+                checkpoint_path = path / metadata['checkpoint']
+                state_dict = torch.load(checkpoint_path, map_location='cpu')
+                model.load_state_dict(state_dict)
+                return model
+            else:
+                raise KeyError("Key 'checkpoint' must be specified.")
+        else:
+            raise ValueError(
+                "ModelType {} or argument 'model={}' is not acceptable for pytorch"
+                " loading.".format(model_type, type(model))
+            )
diff --git a/python/nano/src/bigdl/nano/pytorch/vision/models/_utils.py b/python/nano/src/bigdl/nano/pytorch/vision/models/_utils.py
index e56d37a9e98..b98a76e416a 100644
--- a/python/nano/src/bigdl/nano/pytorch/vision/models/_utils.py
+++ b/python/nano/src/bigdl/nano/pytorch/vision/models/_utils.py
@@ -26,7 +26,6 @@ def set_parameter_requires_grad(model: nn.Module, exception: str):
     :param exception: (str) The name of the layer that does not need to be frozen.
     :return:
     """
-
     for name, param in model.named_parameters():
         if exception not in name:
             param.requires_grad = False
diff --git a/python/nano/src/bigdl/nano/utils/inference/pytorch/model.py b/python/nano/src/bigdl/nano/utils/inference/pytorch/model.py
index 3a6e28a902e..691f743a61d 100644
--- a/python/nano/src/bigdl/nano/utils/inference/pytorch/model.py
+++ b/python/nano/src/bigdl/nano/utils/inference/pytorch/model.py
@@ -15,18 +15,15 @@
 #
 from pytorch_lightning import LightningModule
 import torch
-
+import yaml
+from pathlib import Path
 from .model_utils import get_forward_args
 
 
 class AcceleratedLightningModule(LightningModule):
     def __init__(self, model):
-        self.model = model
         super().__init__()
-        self.on_init_end()
-
-    def on_init_end(self, *args):
-        pass
+        self.model = model
 
     def forward(self, *inputs):
         inputs = self.on_forward_start(inputs)
@@ -62,5 +59,41 @@ def numpy_to_tensors(np_arrays):
             tensors = tensors[0]
         return tensors
 
-    def save(self, path):
-        raise NotImplementedError("AcceleratedLightningModule.save should be overridden.")
+    def _dump_status(self, path):
+        meta_path = Path(path) / "meta-data.yml"
+        with open(meta_path, 'w') as f:
+            yaml.safe_dump(self.status, f)
+
+    def _save_model(self, path):
+        """
+        Save the model file to directory.
+
+        :param path: Path to saved model. Path should be a directory.
+        """
+        raise NotImplementedError("Saving function is not implemented.")
+
+    def _save(self, path):
+        """
+        Save the model to local file.
+
+        :param path: Path to saved model. Path should be a directory.
+        """
+        path = Path(path)
+        Path.mkdir(path, exist_ok=True)
+        self._dump_status(path)
+        self._save_model(path)
+
+    @property
+    def status(self):
+        return {"ModelType": type(self).__name__}
+
+    @staticmethod
+    def _load_status(path):
+        meta_path = Path(path) / "meta-data.yml"
+        with open(meta_path, 'r') as f:
+            metadata = yaml.safe_load(f)
+        return metadata
+
+    @staticmethod
+    def _load(path, model=None):
+        raise NotImplementedError("Loading function is not implemented.")
diff --git a/python/nano/test/openvino/pytorch/test_openvino.py b/python/nano/test/openvino/pytorch/test_openvino.py
index cb914ee9213..a10191bc74b 100644
--- a/python/nano/test/openvino/pytorch/test_openvino.py
+++ b/python/nano/test/openvino/pytorch/test_openvino.py
@@ -14,14 +14,13 @@
 # limitations under the License.
 #
 from unittest import TestCase
-from torchmetrics import F1
 from bigdl.nano.pytorch.trainer import Trainer
 from torchvision.models.mobilenetv3 import mobilenet_v3_small
 import torch
 from torch.utils.data.dataset import TensorDataset
 from torch.utils.data.dataloader import DataLoader
 import os
-
+import shutil
 
 class TestOpenVINO(TestCase):
     def test_trainer_trace_openvino(self):
@@ -47,34 +46,16 @@ def test_trainer_trace_openvino(self):
         y_hat = openvino_model(x)
         assert y_hat.shape == (10, 10)
 
-        openvino_model.save('saved_openvino_model.xml')
-        assert os.path.exists('saved_openvino_model.xml')
-        os.remove('saved_openvino_model.xml')
-
     def test_trainer_save_openvino(self):
         trainer = Trainer(max_epochs=1)
         model = mobilenet_v3_small(num_classes=10)
         x = torch.rand((10, 3, 256, 256))
-        y = torch.ones((10, ), dtype=torch.long)       
 
         # save and load pytorch model
-        Trainer.save(model, 'trainer_save_openvino_model.xml', accelerator='openvino', input_sample=x)
-        assert os.path.exists('trainer_save_openvino_model.xml')
-        openvino_model = Trainer.load('trainer_save_openvino_model.xml')
-        y_hat = openvino_model(x)
-        assert y_hat.shape == (10, 10)
-        os.remove('trainer_save_openvino_model.xml')
-
-        # save and load pytorch-lightning model
-        pl_model = Trainer.compile(model, loss=torch.nn.CrossEntropyLoss(),
-                                   optimizer=torch.optim.SGD(model.parameters(), lr=0.01))
-        ds = TensorDataset(x, y)
-        dataloader = DataLoader(ds, batch_size=2)
-        trainer.fit(pl_model, dataloader)
-
-        Trainer.save(pl_model, 'pl_trainer_save_openvino_model.xml', accelerator='openvino')
-        assert os.path.exists('pl_trainer_save_openvino_model.xml')
-        openvino_model = Trainer.load('pl_trainer_save_openvino_model.xml')
-        y_hat = openvino_model(x)
+        openvino_model = trainer.trace(model, accelerator='openvino', input_sample=x)
+        trainer.save(openvino_model, 'trainer_save_openvino_model')
+        assert len(os.listdir('trainer_save_openvino_model')) > 0
+        loaded_openvino_model = trainer.load('trainer_save_openvino_model')
+        y_hat = loaded_openvino_model(x)
         assert y_hat.shape == (10, 10)
-        os.remove('pl_trainer_save_openvino_model.xml')
+        shutil.rmtree('trainer_save_openvino_model')
diff --git a/python/nano/test/pytorch/tests/test_trainer.py b/python/nano/test/pytorch/tests/test_trainer.py
index 9f70c00cd88..69ea5dff6db 100644
--- a/python/nano/test/pytorch/tests/test_trainer.py
+++ b/python/nano/test/pytorch/tests/test_trainer.py
@@ -15,7 +15,9 @@
 #
 
 
+import copy
 import os
+import shutil
 from unittest import TestCase
 
 import pytest
@@ -76,6 +78,25 @@ def test_trainer_compile(self):
         trainer = Trainer(max_epochs=1)
         pl_model = Trainer.compile(self.model, self.loss, self.optimizer)
         trainer.fit(pl_model, self.train_loader)
+
+    def test_trainer_save_load(self):
+        trainer = Trainer(max_epochs=1)
+        pl_model = Trainer.compile(self.model, self.loss, self.optimizer)
+        trainer.save(pl_model, "saved_model")
+        assert len(os.listdir('saved_model')) > 0
+
+        # save original parameters
+        original_state_dict = copy.deepcopy(pl_model.state_dict())
+
+        # update paramters
+        trainer.fit(pl_model, self.train_loader)
+
+        loaded_pl_model = trainer.load("saved_model", pl_model)
+        loaded_state_dict = loaded_pl_model.state_dict()
+        # check if parameters are updated
+        for k in original_state_dict.keys():
+            assert (original_state_dict[k] == loaded_state_dict[k]).all()
+        shutil.rmtree('saved_model')
         
 
 if __name__ == '__main__':
diff --git a/python/nano/test/run-nano-type-test.sh b/python/nano/test/run-nano-type-test.sh
index e55757f0e35..860e934e75c 100644
--- a/python/nano/test/run-nano-type-test.sh
+++ b/python/nano/test/run-nano-type-test.sh
@@ -3,5 +3,5 @@
 export ANALYTICS_ZOO_ROOT=${ANALYTICS_ZOO_ROOT}
 export NANO_HOME=${ANALYTICS_ZOO_ROOT}/python/nano/
 export NANO_TEST_DIR=${ANALYTICS_ZOO_ROOT}/python/nano/test
-
-mypy --config-file ${NANO_TEST_DIR}/mypy.ini $NANO_HOME/src $NANO_HOME/example
+# "mypy --install-types --non-interactive" is to automatically install missing types 
+mypy --install-types --non-interactive --config-file ${NANO_TEST_DIR}/mypy.ini $NANO_HOME/src $NANO_HOME/example


Code review comments:
diff_hunk: @@ -310,56 +310,74 @@ def trace(model: nn.Module, input_sample=None, accelerator=None):
                              model is a LightningModule with any dataloader attached.
         :param accelerator: The accelerator to use, defaults to None meaning staying in Pytorch
                             backend. Only 'openvino' is supported for now.
-        :return: A PytorchOpenVINOModel using openvino for inference if accelerator='openvino'
+        :return: Model with different acceleration(OpenVINO/ONNX).
         """
         if accelerator == 'openvino':
             return PytorchOpenVINOModel(model, input_sample)
 
     @staticmethod
-    def save(model, path, precision=None, accelerator=None, input_sample=None):
+    def save(model: LightningModule, path):
         """
-        Save the model to path with desired precision and format(by assigning 'accelerator').
-
-        Using a pytorch model, you can export to ONNX, OpenVINO directly with
-        accelerator='openvino'/'onnx'.
+        Save the model to local file.
 
-        :param model: Any model of torch.nn.Module, including PytorchOpenVINOModel
-        :param path: Path to saved model. You need to specify path suffix carefully.
-                     For example, 'model.xml' for OpenVINO, 'model.onnx' for ONNX.
-        :param precision: Precision of saved model, 'FP32', 'FP16', 'INT8', defaults to None,
-                          then the precision depends on the precision of model.
-        :param accelerator: Saved model format and its future accelerator, defaults to None.
-                            accelerator=None will save Pytorch model, accelerator='openvino' will
-                            save an IR to inference by OpenVINO. Same as ONNX.
-        :param input_sample: A set of inputs for trace, defaults to None if you have trace before or
-                             model is a LightningModule with any dataloader attached,
-                             defaults to None.
-                             Only ONNX and OpenVINO requires an input sample to export the model.
+        :param model: Any model of torch.nn.Module, including PytorchOpenVINOModel,
+         PytorchONNXModel.
+        :param path: Path to saved model. Path should be a directory.
         """
-        if accelerator == 'openvino':
-            if precision is None:
-                if not hasattr(model, "save") and isinstance(model, nn.Module):
-                    model = PytorchOpenVINOModel(model, input_sample=input_sample)
-                else:
-                    raise TypeError("Model type of {} can not be exported.".format(type(model)))
-                model.save(path)
+        path = Path(path)
+        Path.mkdir(path, exist_ok=True)
+        if hasattr(model, '_save'):
+            model._save(path)
+        else:
+            # typically for models of nn.Module, LightningModule and LightningModuleFromTorch type
+            meta_path = Path(path) / "meta-data.yml"
path: python/nano/src/bigdl/nano/pytorch/trainer/Trainer.py
text: @user1:
Change to `nano_model_meta.yml`?

@author:
modified.