From 73d52f733ce6259d463fe4a10c92b3e90bab4438 Mon Sep 17 00:00:00 2001 From: Falko Habel Date: Thu, 30 Jan 2025 10:41:37 +0100 Subject: [PATCH] updated config for updated model --- src/aiunn/finetune.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/aiunn/finetune.py b/src/aiunn/finetune.py index 5174d87..ce1b07a 100644 --- a/src/aiunn/finetune.py +++ b/src/aiunn/finetune.py @@ -122,12 +122,17 @@ class ModelTrainer: # Add upscaling layer if not already present if not hasattr(self.model, 'upsample'): + # Get existing configuration values or set defaults if necessary hidden_size = self.model.config.hidden_size kernel_size = self.model.config.kernel_size + self.model.upsample = nn.Sequential( nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False), nn.Conv2d(hidden_size, 3, kernel_size=kernel_size, padding=1) ) + # Update the model's configuration with new parameters + self.model.config.upsample_hidden_size = hidden_size + self.model.config.upsample_kernel_size = kernel_size # Initialize optimizer and loss function self.criterion = nn.MSELoss()