TMP

Cargo.toml

@@ -75,10 +75,12 @@ compact_str = "0.9.0"
 dary_heap = "0.3.8"
 fancy-regex = "0.16.2"
 indicatif = { version = "0.18.3", features = ["improved_unicode"] }
+itertools = "0.14.0"
 futures-util = "0.3.31"
 ndarray = { version = "0.16.1", features = ["serde"] }
 parking_lot = "0.12.5"
 parquet = "56.2.0"
+rand = "0.9.2"
 rayon = "1.11.0"
 reqwest = { version = "0.12.24", features = ["json", "stream"] }
 serde = "1.0.228"
@@ -91,7 +93,10 @@ tracing-indicatif = "0.3.13"
 tracing-subscriber = { version = "0.3.20", features = ["env-filter", "json"] }
 url = "2.5.7"
 
+
+burn-train = { version = "0.19.1", default-features = false }
+
 [workspace.dependencies.burn]
 version = "0.19.1"
 default-features = false
-features = ["std", "fusion", "ndarray", "webgpu", "cuda"]
+features = ["std", "fusion", "ndarray", "webgpu", "cuda", "autodiff"]