updated to the newest version of burn

Cargo.toml

@@ -8,17 +8,13 @@ edition = "2021"
 [features]
 wgpu-backend = ["burn-wgpu"]
 
-[dependencies.burn-wgpu]
-package = "burn-wgpu"
-git = "https://github.com/burn-rs/burn.git"
-optional = true
-
 [dependencies]
-burn = "0.14.0"
-burn-ndarray = "0.14.0"
-burn-tch = "0.14.0"
-burn-autodiff = "0.14.0"
-tch = "0.15.0"
+burn = "0.20.1"
+burn-ndarray = "0.20.1"
+burn-tch = "0.20.1"
+burn-autodiff = "0.20.1"
+burn-wgpu = { version = "0.20.1", optional = true }
+tch = "0.22.0"
 serde = {version = "1.0.171", features = ["std", "derive"]}
 npy = "0.4.0"
 num-traits = "0.2.15"