fix: Fix automatic batch size try catch not handling onnx errors (#123)

* fix: Fix automatic batch size try catch not handling onnx errors * build: Upgrade version, update changelog
orobix · Jun 24, 2024 · 4e07340 · 4e07340
1 parent 44c0a4f
commit 4e07340
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 3 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -2,6 +2,12 @@
 # Changelog
 All notable changes to this project will be documented in this file.
 
+### [2.1.11]
+
+#### Fixed
+
+- Fix sklearn automatic batch finder not working properly with ONNX backbones
+
 ### [2.1.10]
 
 #### Fixed

diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "quadra"
-version = "2.1.10"
+version = "2.1.11"
 description = "Deep Learning experiment orchestration library"
 authors = [
 	"Federico Belotti <[email protected]>",

diff --git a/quadra/__init__.py b/quadra/__init__.py
@@ -1,4 +1,4 @@
-__version__ = "2.1.10"
+__version__ = "2.1.11"
 
 
 def get_version():

diff --git a/quadra/utils/classification.py b/quadra/utils/classification.py
@@ -601,11 +601,12 @@ def automatic_batch_size_computation(
             log.info("Trying batch size: %d", datamodule.batch_size)
             _ = get_feature(feature_extractor=backbone, dl=base_dataloader, iteration_over_training=1, limit_batches=1)
         except RuntimeError as e:
-            if "CUDA out of memory" in str(e):
+            if batch_size > 1:
                 batch_size = batch_size // 2
                 optimal = False
                 continue
 
+            log.error("Unable to run the model with batch size 1")
             raise e
 
         log.info("Found optimal batch size: %d", datamodule.batch_size)