Fix: the new decod

instadeepai · Feb 5, 2025 · b0b3ebe · b0b3ebe
1 parent d81f4cc
commit b0b3ebe
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/instageo/model/model.py b/instageo/model/model.py
@@ -208,9 +208,10 @@ def upscaling_block(in_channels: int, out_channels: int) -> nn.Module:
             )
 
         embed_dims = [
-            (model_args["embed_dim"] * model_args["num_frames"]) // (2**i)
+            (model_args["embed_dim"]) // (2**i)  # Remove num_frames multiplication
             for i in range(5)
         ]
+
         self.segmentation_head = nn.Sequential(
             *[upscaling_block(embed_dims[i], embed_dims[i + 1]) for i in range(4)],
             nn.Conv2d(