add back type annotation that disappeared

modal-labs · Dec 20, 2023 · 27b2482 · 27b2482
1 parent d6f7437
commit 27b2482
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/06_gpu_and_ml/stable_diffusion/stable_diffusion_cli.py b/06_gpu_and_ml/stable_diffusion/stable_diffusion_cli.py
@@ -122,7 +122,9 @@ def __enter__(self):
     __build__ = __enter__
 
     @method()
-    def run_inference(self, prompt: str, steps: int = 20, batch_size: int = 4):
+    def run_inference(
+        self, prompt: str, steps: int = 20, batch_size: int = 4
+    ) -> list[bytes]:
         with torch.inference_mode():
             with torch.autocast("cuda"):
                 images = self.pipe(