diff --git a/crates/burn-dataset/Cargo.toml b/crates/burn-dataset/Cargo.toml
index fd21ec343..d2bf8a229 100644
--- a/crates/burn-dataset/Cargo.toml
+++ b/crates/burn-dataset/Cargo.toml
@@ -21,7 +21,7 @@ fake = ["dep:fake"]
 sqlite = ["__sqlite-shared", "dep:rusqlite"]
 sqlite-bundled = ["__sqlite-shared", "rusqlite/bundled"]
 
-vision = ["dep:flate2", "dep:globwalk", "dep:burn-common"]
+vision = ["dep:flate2", "dep:globwalk", "dep:burn-common", "dep:image"]
 
 # internal
 __sqlite-shared = [