Update docstrings

stefanklut · Mar 25, 2024 · 36460ca · 36460ca
1 parent dfd5cc8
commit 36460ca
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 2 deletions.
diff --git a/api/flask_app.py b/api/flask_app.py
@@ -125,7 +125,14 @@ def setup_model(self, model_name: str, args: DummyArgs):
 
 def safe_predict(data, device):
     """
-    Attempt to predict on the specified device, falling back to CPU on OOM.
+    Attempt to predict on the speci
+
+    Args:
+        data: Data to predict on
+        device: Device to predict on
+
+    Returns:
+        Prediction output
     """
 
     try:

diff --git a/run.py b/run.py
@@ -139,6 +139,7 @@ def cpu_call(self, data: AugInput, device: str = None) -> tuple[dict, int, int]:
 
         Args:
             data (AugInput): image to run the model on
+            device (str): device to run the model on
 
         Returns:
             tuple[dict, int, int]: predictions, height, width
@@ -182,12 +183,13 @@ def cpu_call(self, data: AugInput, device: str = None) -> tuple[dict, int, int]:
 
         return predictions, height, width
 
-    def __call__(self, data: AugInput, device: str = None):
+    def __call__(self, data: AugInput, device: str = None) -> tuple[dict, int, int]:
         """
         Run the model on the image with preprocessing
 
         Args:
             data (AugInput): image to run the model on
+            device (str): device to run the model on
         Returns:
             tuple[dict, int, int]: predictions, height, width
         """