From ef09b637b93315b218ac22a74b304dbc5568577c Mon Sep 17 00:00:00 2001
From: David Chavez <david@dcvz.io>
Date: Tue, 5 Dec 2023 13:30:29 +0100
Subject: [PATCH] chore(web): better panic  (#1051)

---
 Cargo.toml                                   |  1 +
 examples/image-classification-web/Cargo.toml |  1 +
 examples/image-classification-web/src/web.rs | 10 +++++++---
 examples/mnist-inference-web/Cargo.toml      |  3 ++-
 examples/mnist-inference-web/src/web.rs      |  6 ++++++
 5 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 7c6bd22ff..09e6fd4bb 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -74,6 +74,7 @@ wasm-bindgen = "0.2.88"
 wasm-bindgen-futures = "0.4.38"
 wasm-logger = "0.2.0"
 wasm-timer = "0.2.5"
+console_error_panic_hook = "0.1.7"
 
 
 # WGPU stuff
diff --git a/examples/image-classification-web/Cargo.toml b/examples/image-classification-web/Cargo.toml
index 10a4960fb..e4cf7dfff 100644
--- a/examples/image-classification-web/Cargo.toml
+++ b/examples/image-classification-web/Cargo.toml
@@ -29,6 +29,7 @@ wasm-bindgen = { workspace = true }
 wasm-bindgen-futures = { workspace = true }
 wasm-logger = { workspace = true }
 wasm-timer = { workspace = true }
+console_error_panic_hook = { workspace = true }
 
 [build-dependencies]
 # Used to generate code from ONNX model
diff --git a/examples/image-classification-web/src/web.rs b/examples/image-classification-web/src/web.rs
index 18e4bd74b..385afbd76 100644
--- a/examples/image-classification-web/src/web.rs
+++ b/examples/image-classification-web/src/web.rs
@@ -20,6 +20,13 @@ use serde::Serialize;
 use wasm_bindgen::prelude::*;
 use wasm_timer::Instant;
 
+#[wasm_bindgen(start)]
+pub fn start() {
+    // Initialize the logger so that the logs are printed to the console
+    console_error_panic_hook::set_once();
+    wasm_logger::init(wasm_logger::Config::default());
+}
+
 #[allow(clippy::large_enum_variant)]
 /// The model is loaded to a specific backend
 pub enum ModelType {
@@ -49,9 +56,6 @@ impl ImageClassifier {
     /// Constructor called by JavaScripts with the new keyword.
     #[wasm_bindgen(constructor)]
     pub fn new() -> Self {
-        // Initialize the logger so that the logs are printed to the console
-        wasm_logger::init(wasm_logger::Config::default());
-
         log::info!("Initializing the image classifier");
 
         Self {
diff --git a/examples/mnist-inference-web/Cargo.toml b/examples/mnist-inference-web/Cargo.toml
index cf9135928..21b883b17 100644
--- a/examples/mnist-inference-web/Cargo.toml
+++ b/examples/mnist-inference-web/Cargo.toml
@@ -20,7 +20,8 @@ burn = { path = "../../burn", default-features = false }
 serde = { workspace = true }
 wasm-bindgen = { workspace = true }
 wasm-bindgen-futures = { workspace = true }
-js-sys = "0.3.65"
+js-sys = { workspace = true }
+console_error_panic_hook = { workspace = true }
 
 [dev-dependencies]
 pollster = { workspace = true }
diff --git a/examples/mnist-inference-web/src/web.rs b/examples/mnist-inference-web/src/web.rs
index d15aa4a25..d528f0244 100644
--- a/examples/mnist-inference-web/src/web.rs
+++ b/examples/mnist-inference-web/src/web.rs
@@ -11,6 +11,11 @@ use crate::state::{build_and_load_model, Backend};
 
 use burn::tensor::Tensor;
 
+#[cfg_attr(target_family = "wasm", wasm_bindgen(start))]
+pub fn start() {
+    console_error_panic_hook::set_once();
+}
+
 /// Mnist structure that corresponds to JavaScript class.
 /// See:[exporting-rust-struct](https://rustwasm.github.io/wasm-bindgen/contributing/design/exporting-rust-struct.html)
 #[cfg_attr(target_family = "wasm", wasm_bindgen)]
@@ -23,6 +28,7 @@ impl Mnist {
     /// Constructor called by JavaScripts with the new keyword.
     #[cfg_attr(target_family = "wasm", wasm_bindgen(constructor))]
     pub fn new() -> Self {
+        console_error_panic_hook::set_once();
         Self { model: None }
     }