Merge pull request #149 from czbiohub-sf/bump-yogo-ver-to-1

Bump yogo ver to 1
czbiohub-sf · May 31, 2024 · edfcbce · edfcbce
2 parents 86a0f69 + 53c0abc
commit edfcbce
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 3 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "yogo"
-version = "0.2.0"
+version = "1.0.0"
 description = "you only glance once?"
 authors = ["Axel Jacobsen"]
 readme = "README.md"

diff --git a/yogo/infer.py b/yogo/infer.py
@@ -310,9 +310,8 @@ def predict(
 
         # gross! device-type is checked even if enabled=False, which means we
         # have to just tell autocast that device type is always cuda.
-        with torch.autocast(
+        with torch.cuda.amp.autocast(
             enabled=half and device.type == "cuda",
-            device_type="cuda",
             dtype=torch.bfloat16,
         ):
             res = model_jit(img_batch.to(device))