MiT PatchEmbedding input shape fix (#140)

* fix * MiT patch embedding fix * adjustment
DavidLandup0 · Jun 4, 2023 · f45e928 · f45e928
1 parent c256451
commit f45e928
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/deepvision/models/classification/mix_transformer/mit_pt.py b/deepvision/models/classification/mix_transformer/mit_pt.py
@@ -85,7 +85,7 @@ def __init__(
 
         for i in range(self.num_stages):
             patch_embed_layer = OverlappingPatchingAndEmbedding(
-                in_channels=3 if i == 0 else embed_dims[i - 1],
+                in_channels=input_shape[0] if i == 0 else embed_dims[i - 1],
                 out_channels=embed_dims[0] if i == 0 else embed_dims[i],
                 patch_size=7 if i == 0 else 3,
                 stride=4 if i == 0 else 2,

diff --git a/deepvision/models/classification/mix_transformer/mit_tf.py b/deepvision/models/classification/mix_transformer/mit_tf.py
@@ -70,7 +70,7 @@ def __init__(
 
         for i in range(num_stages):
             patch_embed_layer = OverlappingPatchingAndEmbedding(
-                in_channels=3 if i == 0 else embed_dims[i - 1],
+                in_channels=input_shape[-1] if i == 0 else embed_dims[i - 1],
                 out_channels=embed_dims[0] if i == 0 else embed_dims[i],
                 patch_size=7 if i == 0 else 3,
                 stride=4 if i == 0 else 2,