add missing sigmoid to JS

segment-any-text · Sep 21, 2020 · 6169c0c · 6169c0c
1 parent fc6f975
commit 6169c0c
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -127,6 +127,8 @@ jobs:
 
     steps:
       - uses: actions/checkout@v2
+        with:
+          ref: ${{ github.head_ref }}
       - name: Install wasm-pack
         run: curl https://rustwasm.github.io/wasm-pack/installer/init.sh -sSf | sh
       - name: Set up Python 3.8

diff --git a/bindings/javascript/src/tractjs_backend.rs b/bindings/javascript/src/tractjs_backend.rs
@@ -75,9 +75,12 @@ impl TractJSBackend {
         let shape = (shape[0] as usize, shape[1] as usize, shape[2] as usize);
 
         let data: Float32Array = pred.data().into();
-        let preds =
+        let mut preds =
             Array3::from_shape_vec(shape, data.to_vec()).map_err(|_| "Array conversion error")?;
 
+        // sigmoid
+        preds.mapv_inplace(|x| 1f32 / (1f32 + (-x).exp()));
+
         Ok(preds)
     }
 }