Trying to get the CNNTransformer to work, but it is hard.

1 files changed, 5 insertions, 4 deletions

diff --git a/src/text_recognizer/models/base.py b/src/text_recognizer/models/base.py
index cbef787..cc44c92 100644
--- a/src/text_recognizer/models/base.py
+++ b/src/text_recognizer/models/base.py
@@ -141,11 +141,12 @@ class Model(ABC):
             "transform" in self.dataset_args["args"]
             and self.dataset_args["args"]["transform"] is not None
         ):
-            transform_ = [
-                getattr(transforms_module, t["type"])()
-                for t in self.dataset_args["args"]["transform"]
-            ]
+            transform_ = []
+            for t in self.dataset_args["args"]["transform"]:
+                args = t["args"] or {}
+                transform_.append(getattr(transforms_module, t["type"])(**args))
             self.dataset_args["args"]["transform"] = Compose(transform_)
+
         if (
             "target_transform" in self.dataset_args["args"]
             and self.dataset_args["args"]["target_transform"] is not None