diff --git a/src/aiunn/upsampler.py b/src/aiunn/upsampler.py index bc48e97..d14133a 100644 --- a/src/aiunn/upsampler.py +++ b/src/aiunn/upsampler.py @@ -9,7 +9,7 @@ class Upsampler(AIIA): # Assume that base_model.config is an instance of UpsamplerConfig. super().__init__(base_model.config) self.base_model = base_model - self.config = UpsamplerConfig(self.base_model.config) + self.config = UpsamplerConfig(kwargs=self.base_model.config) # Create the upsample layer using values from the configuration. print(self.config.upsample_scale) self.upsample = nn.Upsample(