diff --git a/rust/image-classification/Cargo.lock b/rust/image-classification/Cargo.lock
index 211235fac..f88ef3982 100644
--- a/rust/image-classification/Cargo.lock
+++ b/rust/image-classification/Cargo.lock
@@ -62,6 +62,7 @@ dependencies = [
  "candid",
  "colorgrad",
  "ic-cdk",
+ "ic-stable-structures",
  "ic-wasi-polyfill",
  "image",
  "prost",
diff --git a/rust/image-classification/README.md b/rust/image-classification/README.md
index 0e782b721..939e28d61 100644
--- a/rust/image-classification/README.md
+++ b/rust/image-classification/README.md
@@ -39,6 +39,12 @@ Download MobileNet v2-7 to `src/backend/assets/mobilenetv2-7.onnx`:
 ./downdload_model.sh
 ```
 
+Install NodeJS dependencies for the frontend:
+
+```
+npm install
+```
+
 # Build
 
 ```
diff --git a/rust/image-classification/src/backend/Cargo.toml b/rust/image-classification/src/backend/Cargo.toml
index 3881f2264..048bf833e 100644
--- a/rust/image-classification/src/backend/Cargo.toml
+++ b/rust/image-classification/src/backend/Cargo.toml
@@ -17,4 +17,5 @@ prost-types = "0.11.0"
 image = { version = "0.24", features = ["png"], default-features = false }
 serde = { version = "1.0", features = ["derive"] }
 tract-onnx = { git = "https://github.com/sonos/tract", version = "=0.21.2-pre" }
+ic-stable-structures = "0.6"
 ic-wasi-polyfill = { git = "https://github.com/wasm-forge/ic-wasi-polyfill", version = "0.3.17" }
diff --git a/rust/image-classification/src/backend/src/lib.rs b/rust/image-classification/src/backend/src/lib.rs
index 632d5d82b..04b227d94 100644
--- a/rust/image-classification/src/backend/src/lib.rs
+++ b/rust/image-classification/src/backend/src/lib.rs
@@ -1,7 +1,19 @@
+use std::cell::RefCell;
 use candid::{CandidType, Deserialize};
+use ic_stable_structures::{memory_manager::{MemoryId, MemoryManager}, DefaultMemoryImpl};
 
 mod onnx;
 
+// WASI polyfill requires a virtual stable memory to store the file system.
+// You can replace `0` with any index up to `254`.
+const WASI_MEMORY_ID: MemoryId = MemoryId::new(0);
+
+thread_local! {
+    // The memory manager is used for simulating multiple memories.
+    static MEMORY_MANAGER: RefCell<MemoryManager<DefaultMemoryImpl>> =
+        RefCell::new(MemoryManager::init(DefaultMemoryImpl::default()));
+}
+
 #[derive(CandidType, Deserialize)]
 struct Classification {
     label: String,
@@ -32,11 +44,14 @@ fn classify(image: Vec<u8>) -> ClassificationResult {
 
 #[ic_cdk::init]
 fn init() {
-    ic_wasi_polyfill::init(&[0u8; 32], &[]);
+    let wasi_memory = MEMORY_MANAGER.with(|m| m.borrow().get(WASI_MEMORY_ID));
+    ic_wasi_polyfill::init_with_memory(&[0u8; 32], &[], wasi_memory);
     onnx::setup().unwrap();
 }
 
 #[ic_cdk::post_upgrade]
 fn post_upgrade() {
+    let wasi_memory = MEMORY_MANAGER.with(|m| m.borrow().get(WASI_MEMORY_ID));
+    ic_wasi_polyfill::init_with_memory(&[0u8; 32], &[], wasi_memory);
     onnx::setup().unwrap();
 }