Merge 79efd27a7e into 17a886cb58

docker/serve/Dockerfile

@@ -1,6 +1,7 @@
 ARG PYTORCH="2.0.1"
 ARG CUDA="11.7"
 ARG CUDNN="8"
+ARG TORCHSERVE="0.7.1"
 FROM pytorch/torchserve:latest-gpu
 
 ARG MMPRE="1.2.0"
@@ -11,8 +12,8 @@ ENV HOME="/home/model-server"
 ENV PATH="/opt/conda/bin:$HOME/.local/bin:$PATH"
 RUN export FORCE_CUDA=1
 
-# TORCHSEVER
-RUN pip install torchserve torch-model-archiver
+# TORCHSERVE
+RUN pip install torchserve==${TORCHSERVE} torch-model-archiver==${TORCHSERVE}
 RUN pip install nvgpu
 
 # OPEN-MMLAB

tools/torchserve/mmpretrain2torchserve.py

@@ -10,7 +10,8 @@ try:
     from model_archiver.model_packaging_utils import ModelExportUtils
 except ImportError:
     raise ImportError(
-        'Please run `pip install torchserve torch-model-archiver"` to '
+        'Please run '
+        '`pip install torchserve==0.7.1 torch-model-archiver==0.7.1"` to '
         'install required third-party libraries.')
 
 
@@ -106,7 +107,7 @@ if __name__ == '__main__':
 
     if package_model is None:
         raise ImportError('`torch-model-archiver` is required.'
-                          'Try: pip install torch-model-archiver')
+                          'Try: pip install torch-model-archiver==0.7.1')
 
     mmpretrain2torchserve(args.config, args.checkpoint, args.output_folder,
                           args.model_name, args.model_version, args.force)