1 files changed, 3 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 4644003..cdbfd74 100644
--- a/workspaces/main/shaders/cnn_v2/cnn_v2_compute.wgsl
+++ b/workspaces/main/shaders/cnn_v2/cnn_v2_compute.wgsl
@@ -122,12 +122,10 @@ fn main(@builtin(global_invocation_id) id: vec3<u32>) {
     }
 
     // Activation (matches train_cnn_v2.py)
-    if (is_output) {
-      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]
+    if (is_output || params.is_layer_0 != 0u) {
+      output[c] = 1.0 / (1.0 + exp(-sum));  // Sigmoid [0,1]
     } else {
-      output[c] = max(0.0, sum);  // Middle layers: ReLU
+      output[c] = max(0.0, sum);  // ReLU
     }
   }