For ViTs, Image size should be extracted from their conf

naamiinepal · Dec 30, 2023 · 81a9a58 · 81a9a58
1 parent 69e7eed
commit 81a9a58
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 7 deletions.
diff --git a/configs/model/trans_seg.yaml b/configs/model/trans_seg.yaml
@@ -13,7 +13,6 @@ net:
     norm_first: true
   num_decoder_layers: 4
   num_upsampler_layers: 5
-  final_image_size: ${img_size}
   num_output_channels: 1
 
 loss_fn:

diff --git a/src/models/core_models/trans_segmentor/__init__.py b/src/models/core_models/trans_segmentor/__init__.py
@@ -1,7 +1,7 @@
 # pyright: reportGeneralTypeIssues=false
 from functools import lru_cache
 from pathlib import Path
-from typing import Mapping, Optional, Union
+from typing import Mapping, Union
 
 import torch
 from torch import nn
@@ -25,7 +25,6 @@ def __init__(
         decoder_layer_kwargs: StrToAny,
         num_decoder_layers: int,
         num_upsampler_layers: int,
-        final_image_size: Optional[int] = None,
         num_output_channels: int = 1,
         *args,
         **kwargs,
@@ -63,16 +62,13 @@ def __init__(
             image_hidden_size=image_hidden_size,
         )
 
-        # Use the image size from the image encoder, if not provided explicitly
-        image_size = final_image_size or img_config.image_size
-
         self.decoder = TransDecoder(
             image_hidden_size=image_hidden_size,
             decoder_layer_kwargs=decoder_layer_kwargs,
             num_decoder_layers=num_decoder_layers,
             patch_size=img_config.patch_size,
             num_upsampler_layers=num_upsampler_layers,
-            final_image_size=image_size,
+            final_image_size=img_config.image_size,
             num_output_channels=num_output_channels,
         )