update examples/pdf_document_extraction/images.py

tensorlakeai · Oct 19, 2024 · a42842c · a42842c
1 parent 4766072
commit a42842c
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/examples/pdf_document_extraction/images.py b/examples/pdf_document_extraction/images.py
@@ -4,8 +4,9 @@
 chroma_image = Image(python="3.11").name("tensorlake/blueprints-chromadb").run("pip install chromadb").run("pip install pillow")
 
 st_image = (
-    Image(python="3.11")
+    Image()
     .name("tensorlake/pdf-blueprint-st")
+    .base_image("pytorch/pytorch:2.4.1-cuda11.8-cudnn9-runtime")
     .run("pip install sentence-transformers")
     .run("pip install langchain")
     .run("pip install pillow")