upgrading to 0.14.0

Cargo.toml

@@ -18,11 +18,11 @@ burn = { git = "https://github.com/burn-rs/burn.git" }
 burn-ndarray = { package = "burn-ndarray", git = "https://github.com/burn-rs/burn.git" }
 burn-tch = { package = "burn-tch", git = "https://github.com/burn-rs/burn.git" }
 burn-autodiff = { package = "burn-autodiff", git = "https://github.com/burn-rs/burn.git" }
-tch = "0.13.0"
-serde = {version = "1.0.171", features = ["std", "derive"]}
+tch = "0.14.0"
+serde = {version = "1.0.192", features = ["std", "derive"]}
 npy = "0.4.0"
 num-traits = "0.2.15"
 rust_tokenizers = "8.1.0"
 regex = "1.9.1"
 image = "0.24.6"
-cfg-if = "0.1"
+cfg-if = "0.1.0"

src/backend.rs

@@ -23,12 +23,13 @@ pub trait Backend: burn::tensor::backend::Backend {
     }
 }
 
+use burn::backend;
 use burn::tensor::ops::TensorOps;
 use burn::tensor::Float;
 use burn_tch::{self, TchElement, TchTensor};
 use tch;
 
-impl<E: TchElement> Backend for burn_tch::TchBackend<E> {
+impl<E: TchElement> Backend for burn_tch::LibTorch<E> {
     fn qkv_attention(
         q: Self::TensorPrimitive<3>,
         k: Self::TensorPrimitive<3>,
@@ -68,18 +69,14 @@ impl<E: TchElement> Backend for burn_tch::TchBackend<E> {
                 Some(mask.tensor),
                 0.0,
                 false,
+                None,
             ),
         ))
         .swap_dims(1, 2)
         .flatten(2, 3)
         .into_primitive()
     }
 }
-
-use burn_autodiff;
-
-impl<B: Backend> Backend for burn_autodiff::ADBackendDecorator<B> {}
-
 use std::f32::NEG_INFINITY;
 
 fn qkv_attention<B: Backend>(

src/bin/convert/main.rs

@@ -11,7 +11,7 @@ use burn::{
     tensor::{backend::Backend, Tensor},
 };
 
-use burn_ndarray::{NdArrayBackend, NdArrayDevice};
+use burn_ndarray::{NdArray, NdArrayDevice};
 
 use burn::record::{self, BinFileRecorder, FullPrecisionSettings, Recorder};
 
@@ -37,7 +37,7 @@ fn save_model_file<B: Backend>(
 }
 
 fn main() {
-    type Backend = NdArrayBackend<f32>;
+    type Backend = NdArray<f32>;
     let device = NdArrayDevice::Cpu;
 
     let args: Vec<String> = env::args().collect();

src/bin/sample/main.rs

@@ -14,16 +14,15 @@ cfg_if::cfg_if! {
     if #[cfg(feature = "wgpu-backend")] {
         use burn_wgpu::{WgpuBackend, WgpuDevice, AutoGraphicsApi};
     } else {
-        use burn_tch::{TchBackend, TchDevice};
+        use burn_tch::{LibTorch, LibTorchDevice};
     }
 }
 
+use burn::record::{self, BinFileRecorder, FullPrecisionSettings, Recorder};
 use std::env;
 use std::io;
 use std::process;
 
-use burn::record::{self, BinFileRecorder, FullPrecisionSettings, Recorder};
-
 fn load_stable_diffusion_model_file<B: Backend>(
     filename: &str,
 ) -> Result<StableDiffusion<B>, record::RecorderError> {
@@ -38,8 +37,8 @@ fn main() {
             type Backend = WgpuBackend<AutoGraphicsApi, f32, i32>;
             let device = WgpuDevice::BestAvailable;
         } else {
-            type Backend = TchBackend<f32>;
-            let device = TchDevice::Cuda(0);
+            type Backend = LibTorch<f32>;
+            let device = LibTorchDevice::Cuda(0);
         }
     }