fix(Architectures): Fixed incorrect input shapes for B16 and B32 Visi…

…on Transformer Architecture
frankkramer-lab · May 28, 2022 · e44916b · e44916b
1 parent 33068ac
commit e44916b
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/aucmedi/neural_network/architectures/image/vit_b16.py b/aucmedi/neural_network/architectures/image/vit_b16.py
@@ -93,7 +93,7 @@ class Architecture_ViT_B16(Architecture_Base):
     #---------------------------------------------#
     #                Initialization               #
     #---------------------------------------------#
-    def __init__(self, classification_head, channels, input_shape=(384, 384),
+    def __init__(self, classification_head, channels, input_shape=(224, 224),
                  pretrained_weights=False):
         self.classifier = classification_head
         self.input = input_shape + (channels,)

diff --git a/aucmedi/neural_network/architectures/image/vit_b32.py b/aucmedi/neural_network/architectures/image/vit_b32.py
@@ -93,7 +93,7 @@ class Architecture_ViT_B32(Architecture_Base):
     #---------------------------------------------#
     #                Initialization               #
     #---------------------------------------------#
-    def __init__(self, classification_head, channels, input_shape=(384, 384),
+    def __init__(self, classification_head, channels, input_shape=(224, 224),
                  pretrained_weights=False):
         self.classifier = classification_head
         self.input = input_shape + (channels,)