From 7ef047460fc5fe3ce84445c93d2467681479e065 Mon Sep 17 00:00:00 2001
From: Guillaume Lagrange <lagrange.guillaume.1@gmail.com>
Date: Fri, 13 Sep 2024 15:01:05 -0400
Subject: [PATCH] Update to burn 0.14

---
 resnet-burn/Cargo.toml                               |  4 ++--
 resnet-burn/examples/finetune/src/data.rs            | 10 +++++-----
 resnet-burn/examples/finetune/src/inference.rs       |  3 +--
 resnet-burn/examples/inference/examples/inference.rs |  8 +++-----
 resnet-burn/examples/inference/src/imagenet.rs       |  4 ++--
 5 files changed, 13 insertions(+), 16 deletions(-)

diff --git a/resnet-burn/Cargo.toml b/resnet-burn/Cargo.toml
index 18c5194..e102010 100644
--- a/resnet-burn/Cargo.toml
+++ b/resnet-burn/Cargo.toml
@@ -17,8 +17,8 @@ license = "MIT OR Apache-2.0"
 
 [workspace.dependencies]
 # Note: default-features = false is needed to disable std
-burn = { version = "0.13.0", default-features = false }
-burn-import = "0.13.0"
+burn = { version = "0.14.0", default-features = false }
+burn-import = "0.14.0"
 dirs = "5.0.1"
 serde = { version = "1.0.192", default-features = false, features = [
     "derive",
diff --git a/resnet-burn/examples/finetune/src/data.rs b/resnet-burn/examples/finetune/src/data.rs
index ce701e3..62cf813 100644
--- a/resnet-burn/examples/finetune/src/data.rs
+++ b/resnet-burn/examples/finetune/src/data.rs
@@ -55,8 +55,8 @@ pub struct Normalizer<B: Backend> {
 impl<B: Backend> Normalizer<B> {
     /// Creates a new normalizer.
     pub fn new(device: &Device<B>) -> Self {
-        let mean = Tensor::from_floats(MEAN, device).reshape([1, 3, 1, 1]);
-        let std = Tensor::from_floats(STD, device).reshape([1, 3, 1, 1]);
+        let mean = Tensor::<B, 1>::from_floats(MEAN, device).reshape([1, 3, 1, 1]);
+        let std = Tensor::<B, 1>::from_floats(STD, device).reshape([1, 3, 1, 1]);
         Self { mean, std }
     }
 
@@ -117,9 +117,9 @@ impl<B: Backend> Batcher<ImageDatasetItem, ClassificationBatch<B>> for Classific
 
         let images = items
             .into_iter()
-            .map(|item| Data::new(image_as_vec_u8(item), Shape::new([HEIGHT, WIDTH, 3])))
-            .map(|data| Tensor::<B, 3>::from_data(data.convert(), &self.device).permute([2, 0, 1]))
-            .map(|tensor| tensor / 255) // normalize between [0, 1]
+            .map(|item| TensorData::new(image_as_vec_u8(item), Shape::new([HEIGHT, WIDTH, 3])))
+            .map(|data| Tensor::<B, 3>::from_data(data.convert::<B::FloatElem>(), &self.device))
+            .map(|tensor| tensor.permute([2, 0, 1]) / 255) // normalize between [0, 1]
             .collect();
 
         let images = Tensor::stack(images, 0);
diff --git a/resnet-burn/examples/finetune/src/inference.rs b/resnet-burn/examples/finetune/src/inference.rs
index aa0defb..a1312ad 100644
--- a/resnet-burn/examples/finetune/src/inference.rs
+++ b/resnet-burn/examples/finetune/src/inference.rs
@@ -46,8 +46,7 @@ pub fn infer<B: Backend>(artifact_dir: &str, device: B::Device, threshold: f32)
     // Get predicted class names over the specified threshold
     let predicted = output.greater_equal_elem(threshold).nonzero()[1]
         .to_data()
-        .value
-        .iter()
+        .iter::<B::IntElem>()
         .map(|i| CLASSES[i.elem::<i64>() as usize])
         .collect::<Vec<_>>();
 
diff --git a/resnet-burn/examples/inference/examples/inference.rs b/resnet-burn/examples/inference/examples/inference.rs
index a26e597..1554f8d 100644
--- a/resnet-burn/examples/inference/examples/inference.rs
+++ b/resnet-burn/examples/inference/examples/inference.rs
@@ -5,7 +5,7 @@ use burn::{
     backend::NdArray,
     module::Module,
     record::{FullPrecisionSettings, NamedMpkFileRecorder},
-    tensor::{backend::Backend, Data, Device, Element, Shape, Tensor},
+    tensor::{backend::Backend, Device, Element, Tensor, TensorData},
 };
 
 const MODEL_PATH: &str = "resnet18-ImageNet1k";
@@ -17,10 +17,8 @@ fn to_tensor<B: Backend, T: Element>(
     shape: [usize; 3],
     device: &Device<B>,
 ) -> Tensor<B, 3> {
-    Tensor::<B, 3>::from_data(Data::new(data, Shape::new(shape)).convert(), device)
-        // permute(2, 0, 1)
-        .swap_dims(2, 1) // [H, C, W]
-        .swap_dims(1, 0) // [C, H, W]
+    Tensor::<B, 3>::from_data(TensorData::new(data, shape).convert::<B::FloatElem>(), device)
+        .permute([2, 0, 1]) // [C, H, W]
         / 255 // normalize between [0, 1]
 }
 
diff --git a/resnet-burn/examples/inference/src/imagenet.rs b/resnet-burn/examples/inference/src/imagenet.rs
index e1dbf1c..075a2f9 100644
--- a/resnet-burn/examples/inference/src/imagenet.rs
+++ b/resnet-burn/examples/inference/src/imagenet.rs
@@ -21,8 +21,8 @@ pub struct Normalizer<B: Backend> {
 impl<B: Backend> Normalizer<B> {
     /// Creates a new normalizer.
     pub fn new(device: &Device<B>) -> Self {
-        let mean = Tensor::from_floats(MEAN, device).reshape([1, 3, 1, 1]);
-        let std = Tensor::from_floats(STD, device).reshape([1, 3, 1, 1]);
+        let mean = Tensor::<B, 1>::from_floats(MEAN, device).reshape([1, 3, 1, 1]);
+        let std = Tensor::<B, 1>::from_floats(STD, device).reshape([1, 3, 1, 1]);
         Self { mean, std }
     }