src/aiunn/finetune.py

@@ -274,6 +274,8 @@ class FineTuner:
             # Forward pass (we'll use the model's existing architecture without adding layers)
             try:
                 features = self.model(low_ress)
+                print("Features shape:", features.shape)  # Check output dimensions
+                print("High-res shape:", high_ress.shape)  # Check target dimensions
             except Exception as e:
                 raise RuntimeError(f"Error during forward pass: {str(e)}")