fix: skip safetensor patching if not installed (#5025)

* fix: skip safetensor patching if not installed Signed-off-by: Frost Ming <[email protected]>
bentoml · Oct 16, 2024 · b41bf29 · b41bf29
1 parent f8874ef
commit b41bf29
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/src/_bentoml_impl/worker/service.py b/src/_bentoml_impl/worker/service.py
@@ -14,7 +14,11 @@ def patch_safetensor():
     import os
     import subprocess
 
-    import safetensors.torch
+    try:
+        import safetensors.torch
+    except ImportError:
+        logger.info("safetensors not installed, skipping model preheat")
+        return
 
     logger.info(
         "Patching safetensors.torch.safe_open to preheat model loading in parallel"