prevent bad things to happen.

wavegru_mod.cc

@@ -122,18 +122,18 @@ struct WaveGRU {
       }
       o1.SpMM_bias(h, o1b, &fco1, true);
       o2.SpMM_bias(fco1, o2b, &fco2, false);
-      // auto max_logit = fco2[0];
-      // for (int i = 1; i <= 255; ++i) {
-      //   max_logit = max(max_logit, fco2[i]);
-      // }
-      // float total = 0.0;
-      // for (int i = 0; i <= 255; ++i) {
-      //   logits[i] = csrblocksparse::fast_exp(fco2[i] - max_logit);
-      //   total += logits[i];
-      // }
-      // for (int i = 0; i <= 255; ++i) {
-      //   if (logits[i] < total / 256.0) fco2[i] = -1e9;
-      // }
+      auto max_logit = fco2[0];
+      for (int i = 1; i <= 255; ++i) {
+        max_logit = max(max_logit, fco2[i]);
+      }
+      float total = 0.0;
+      for (int i = 0; i <= 255; ++i) {
+        logits[i] = csrblocksparse::fast_exp(fco2[i] - max_logit);
+        total += logits[i];
+      }
+      for (int i = 0; i <= 255; ++i) {
+        if (logits[i] < total / 256.0) fco2[i] = -1e9;
+      }
       value = fco2.Sample(temperature);
       signal[index] = value;
     }

wavegru_mod.so

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:700f2cade76db615b1e38bddfc9c604ff1c8ea1af3e507f879d0ceebae5d232d
+oid sha256:b65c5312f24f8ab9cfa51e8340a24ac1165b247046a331386d636fba9036c19c
 size 525536