UPDATE

Dockerfile

@@ -50,8 +50,7 @@ ENV HOME=/home/user \
     PATH=/usr/local/cuda/bin:/home/user/.local/bin:$PATH \
     LD_LIBRARY_PATH="/usr/local/cuda/lib64:${LD_LIBRARY_PATH}" \
     CXX=/usr/bin/g++ \
-    TORCH_CUDA_ARCH_LIST="7.5" \
-    PYTORCH_CUDA_ALLOC_CONF="expandable_segments:True"
+    TORCH_CUDA_ARCH_LIST="7.5"
 WORKDIR $HOME/app
 
 COPY --chown=user . $HOME/app

app.py

@@ -5,7 +5,7 @@ if os.environ.get("MODELSCOPE_ENVIRONMENT") == "studio":
 
     patch_hub()
 
-os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:512"
 
 
 from config import CONFIG, ModelConfig
@@ -135,7 +135,7 @@ for model_config in CONFIG.MODELS:
         MODEL_STORAGE[model_config.SERVICE_NAME].model, model_config.VOCAB
     )
     if "cuda" in CONFIG.STRATEGY:
-        # torch.cuda.empty_cache()
+        torch.cuda.empty_cache()
         gc.collect()
     logGPUState()