2 files changed, 8 insertions, 5 deletions
diff --git a/workspaces/main/shaders/cnn_v2/cnn_v2_compute.wgsl b/workspaces/main/shaders/cnn_v2/cnn_v2_compute.wgsl
index 6905e75..4644003 100644
--- a/workspaces/main/shaders/cnn_v2/cnn_v2_compute.wgsl
+++ b/workspaces/main/shaders/cnn_v2/cnn_v2_compute.wgsl
@@ -11,6 +11,7 @@ struct LayerParams {
   weight_offset: u32,     // Offset in f16 units
   is_output_layer: u32,   // 1 if final layer (sigmoid), 0 otherwise (relu)
   blend_amount: f32,      // [0,1] blend with original
+  is_layer_0: u32,        // 1 if first layer (clamp [0,1]), 0 otherwise
 }
 
 @group(0) @binding(0) var static_features: texture_2d<u32>;       // 8D static features (p0-p3 + spatial)
@@ -120,11 +121,13 @@ fn main(@builtin(global_invocation_id) id: vec3<u32>) {
       }
     }
 
-    // Activation
+    // Activation (matches train_cnn_v2.py)
     if (is_output) {
-      output[c] = clamp(sum, 0.0, 1.0);
+      output[c] = clamp(sum, 0.0, 1.0);  // Output layer: clamp [0,1]
+    } else if (params.is_layer_0 != 0u) {
+      output[c] = clamp(sum, 0.0, 1.0);  // Layer 0: clamp [0,1]
     } else {
-      output[c] = max(0.0, sum);  // ReLU
+      output[c] = max(0.0, sum);  // Middle layers: ReLU
     }
   }
 
diff --git a/workspaces/main/shaders/cnn_v2/cnn_v2_static.wgsl b/workspaces/main/shaders/cnn_v2/cnn_v2_static.wgsl
index 35068a2..7b08132 100644
--- a/workspaces/main/shaders/cnn_v2/cnn_v2_static.wgsl
+++ b/workspaces/main/shaders/cnn_v2/cnn_v2_static.wgsl
@@ -48,9 +48,9 @@ fn main(@builtin(global_invocation_id) id: vec3<u32>) {
   let p2 = rgba.b;
   let p3 = textureLoad(depth_tex, coord, 0).r;
 
-  // UV coordinates (normalized [0,1], bottom-left origin)
+  // UV coordinates (normalized [0,1], top-left origin - matches training)
   let uv_x = f32(coord.x) / f32(dims.x);
-  let uv_y = 1.0 - (f32(coord.y) / f32(dims.y));
+  let uv_y = f32(coord.y) / f32(dims.y);
 
   // Multi-frequency position encoding
   let sin20_y = sin(20.0 * uv_y);