summaryrefslogtreecommitdiff
path: root/src/text_recognizer/networks/util.py
diff options
context:
space:
mode:
authoraktersnurra <gustaf.rydholm@gmail.com>2020-11-08 12:41:04 +0100
committeraktersnurra <gustaf.rydholm@gmail.com>2020-11-08 12:41:04 +0100
commitbeeaef529e7c893a3475fe27edc880e283373725 (patch)
tree59eb72562bf7a5a9470c2586e6280600ad94f1ae /src/text_recognizer/networks/util.py
parent4d7713746eb936832e84852e90292936b933e87d (diff)
Trying to get the CNNTransformer to work, but it is hard.
Diffstat (limited to 'src/text_recognizer/networks/util.py')
-rw-r--r--src/text_recognizer/networks/util.py4
1 files changed, 2 insertions, 2 deletions
diff --git a/src/text_recognizer/networks/util.py b/src/text_recognizer/networks/util.py
index 0d08506..b31e640 100644
--- a/src/text_recognizer/networks/util.py
+++ b/src/text_recognizer/networks/util.py
@@ -28,7 +28,7 @@ def sliding_window(
c = images.shape[1]
patches = unfold(images)
patches = rearrange(
- patches, "b (c h w) t -> b t c h w", c=c, h=patch_size[0], w=patch_size[1]
+ patches, "b (c h w) t -> b t c h w", c=c, h=patch_size[0], w=patch_size[1],
)
return patches
@@ -77,7 +77,7 @@ def configure_backbone(backbone: str, backbone_args: Dict) -> Type[nn.Module]:
if "remove_layers" in backbone_args and backbone_args["remove_layers"] is not None:
backbone = nn.Sequential(
- *list(backbone.children())[0][: -backbone_args["remove_layers"]]
+ *list(backbone.children())[:][: -backbone_args["remove_layers"]]
)
return backbone