diff --git a/src/aiunn/finetune.py b/src/aiunn/finetune.py
index 36ad40a..50acea1 100644
--- a/src/aiunn/finetune.py
+++ b/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)}")