1 files changed, 6 insertions, 3 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
     }
   }