diff --git a/Cargo.lock b/Cargo.lock
index e639b2e..3a27350 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2163,7 +2163,7 @@ checksum = "e2db585e1d738fc771bf08a151420d3ed193d9d895a36df7f6f8a9456b911ddc"
 
 [[package]]
 name = "kord"
-version = "0.7.0"
+version = "0.7.1"
 dependencies = [
  "anyhow",
  "async-trait",
diff --git a/Cargo.toml b/Cargo.toml
index 6b3e5e0..b1d08b3 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "kord"
-version = "0.7.0"
+version = "0.7.1"
 edition = "2021"
 authors = ["Aaron Roney <twitchax@gmail.com>"]
 license = "MIT"
@@ -47,11 +47,13 @@ analyze_file_alac = ["symphonia/alac", "symphonia/isomp4"]
 ml = ["ml_train", "ml_infer"]
 ml_base = ["serde", "byteorder", "bincode"]
 ml_train = ["ml_base", "rand", "rayon", "burn-autodiff", "burn/train", "burn/std", "burn/wgpu", "burn/metrics", "burn-ndarray/std", "burn/tui"]
-ml_infer = ["ml_base", "burn", "burn-ndarray"]
+ml_infer = ["ml_base", "burn", "burn-ndarray", "burn-ndarray/std"]
 ml_gpu = ["ml_train", "burn-tch", "burn-wgpu"]
 
 wasm = ["rodio/wasm-bindgen", "wasm-bindgen", "wasm-bindgen-futures", "js-sys", "console_error_panic_hook", "wee_alloc", "gloo-timers", "burn/wasm-sync"]
 
+wasi = ["burn/wasm-sync"]
+
 plot = ["plotters"]
 
 [dependencies]
diff --git a/DEVELOPMENT.md b/DEVELOPMENT.md
index 8281140..2c2b99d 100644
--- a/DEVELOPMENT.md
+++ b/DEVELOPMENT.md
@@ -31,6 +31,6 @@ $ wasm-pack publish
 ## Publish to wasmer
 
 ```bash
-$ cargo wasi build --release --no-default-features --features cli --features ml_infer --features analyze_file
+$ cargo wasi build --release --no-default-features --features wasi --features cli --features ml_infer --features analyze_file
 $ cargo wasix publish
 ```
\ No newline at end of file
diff --git a/model/model_config.json b/model/model_config.json
index 2dc9242..491b786 100644
--- a/model/model_config.json
+++ b/model/model_config.json
@@ -6,9 +6,8 @@
   "simulation_peak_radius": 2.0,
   "simulation_harmonic_decay": 0.1,
   "simulation_frequency_wobble": 0.4,
-  "mlp_layers": 3,
-  "mlp_size": 1024,
-  "mlp_dropout": 0.3,
+  "mha_heads": 8,
+  "mha_dropout": 0.3,
   "model_epochs": 64,
   "model_batch_size": 100,
   "model_workers": 64,
diff --git a/model/state.json.bin b/model/state.json.bin
index 881ec62..9edd899 100644
Binary files a/model/state.json.bin and b/model/state.json.bin differ
diff --git a/rust-toolchain.toml b/rust-toolchain.toml
deleted file mode 100644
index d3b95a8..0000000
--- a/rust-toolchain.toml
+++ /dev/null
@@ -1,2 +0,0 @@
-[toolchain]
-channel = "nightly-2023-12-10"
\ No newline at end of file
diff --git a/src/core/note.rs b/src/core/note.rs
index 686c243..6dd1dd8 100644
--- a/src/core/note.rs
+++ b/src/core/note.rs
@@ -507,7 +507,7 @@ impl PartialOrd for Note {
 
 impl Ord for Note {
     fn cmp(&self, other: &Self) -> Ordering {
-        self.partial_cmp(other).unwrap_or(Ordering::Equal)
+        self.frequency().partial_cmp(&other.frequency()).unwrap_or(Ordering::Equal)
     }
 }