Move model to device using from_pretrained

CHANGELOG.md

@@ -1,3 +1,7 @@
+## 1.0.1
+
+* fix: moving the table transformer model to device when loading the model instead of once the model is loaded.
+
 ## 1.0.0
 
 * feat: support for Python 3.10+; drop support for Python 3.9

requirements/base.in

@@ -10,6 +10,7 @@ torch
 timm
 # NOTE(alan): Pinned because this is when the most recent module we import appeared
 transformers>=4.25.1
+accelerate
 rapidfuzz
 pandas
 scipy

requirements/base.txt

@@ -4,6 +4,8 @@
 #
 #    pip-compile requirements/base.in
 #
+accelerate==1.7.0
+    # via -r requirements/base.in
 certifi==2025.4.26
     # via requests
 cffi==1.17.1
@@ -36,6 +38,7 @@ fsspec==2025.3.2
 huggingface-hub==0.31.2
     # via
     #   -r requirements/base.in
+    #   accelerate
     #   timm
     #   tokenizers
     #   transformers
@@ -58,6 +61,7 @@ networkx==3.4.2
 numpy==2.2.5
     # via
     #   -r requirements/base.in
+    #   accelerate
     #   contourpy
     #   matplotlib
     #   onnx
@@ -75,6 +79,7 @@ opencv-python==4.11.0.86
     # via -r requirements/base.in
 packaging==25.0
     # via
+    #   accelerate
     #   huggingface-hub
     #   matplotlib
     #   onnxruntime
@@ -91,6 +96,8 @@ protobuf==6.31.0
     # via
     #   onnx
     #   onnxruntime
+psutil==7.0.0
+    # via accelerate
 pycparser==2.22
     # via cffi
 pyparsing==3.2.3
@@ -107,6 +114,7 @@ pytz==2025.2
     # via pandas
 pyyaml==6.0.2
     # via
+    #   accelerate
     #   huggingface-hub
     #   timm
     #   transformers
@@ -120,6 +128,7 @@ requests==2.32.3
     #   transformers
 safetensors==0.5.3
     # via
+    #   accelerate
     #   timm
     #   transformers
 scipy==1.15.3
@@ -137,6 +146,7 @@ tokenizers==0.21.1
 torch==2.7.0
     # via
     #   -r requirements/base.in
+    #   accelerate
     #   timm
     #   torchvision
 torchvision==0.22.0

requirements/dev.txt

@@ -264,7 +264,9 @@ prompt-toolkit==3.0.51
     #   ipython
     #   jupyter-console
 psutil==7.0.0
-    # via ipykernel
+    # via
+    #   -c requirements/base.txt
+    #   ipykernel
 ptyprocess==0.7.0
     # via
     #   pexpect
@@ -351,7 +353,7 @@ terminado==0.18.1
     #   jupyter-server-terminals
 tinycss2==1.4.0
     # via bleach
-tornado==6.4.2
+tornado==6.5
     # via
     #   ipykernel
     #   jupyter-client

unstructured_inference/__version__.py

@@ -1 +1 @@
-__version__ = "1.0.0"  # pragma: no cover
+__version__ = "1.0.1"  # pragma: no cover

unstructured_inference/models/tables.py

@@ -67,7 +67,9 @@ def initialize(
             logger.info("Loading the table structure model ...")
             cached_current_verbosity = logging.get_verbosity()
             logging.set_verbosity_error()
-            self.model = TableTransformerForObjectDetection.from_pretrained(model)
+            self.model = TableTransformerForObjectDetection.from_pretrained(
+                model, device_map=self.device
+            )
             logging.set_verbosity(cached_current_verbosity)
             self.model.eval()
 
@@ -77,7 +79,6 @@ def initialize(
             raise ImportError(
                 "Review the parameters to initialize a UnstructuredTableTransformerModel obj",
             )
-        self.model.to(device)
 
     def get_structure(
         self,