Fix problems with batch norm on LibTorch backend (#1226)

* Fix problems with batch norm

* Add comment

burn-core/src/nn/norm/batch.rs

@@ -106,6 +106,7 @@ impl<const D: usize, B: Backend> BatchNorm<B, D> {
     }
 
     fn forward_train<const DI: usize>(&self, input: Tensor<B, DI>) -> Tensor<B, DI> {
+        let device = input.device();
         let dims = input.dims();
         let batch_size = dims[0];
         let channels = dims[1];
@@ -134,8 +135,8 @@ impl<const D: usize, B: Backend> BatchNorm<B, D> {
             .mean_dim(1)
             .reshape(shape_unsqueeze);
 
-        let running_mean = self.running_mean.value_sync().to_device(&mean.device());
-        let running_var = self.running_var.value_sync().to_device(&var.device());
+        let running_mean = self.running_mean.value_sync().to_device(&device);
+        let running_var = self.running_var.value_sync().to_device(&device);
 
         let running_mean = running_mean.mul_scalar(1.0 - self.momentum).add(
             mean.clone()

burn-tch/src/ops/base.rs

@@ -1,14 +1,29 @@
 use burn_tensor::Shape;
 use tch::Scalar;
 
-use crate::{TchShape, TchTensor};
+use crate::{LibTorchDevice, TchShape, TchTensor};
 use std::{marker::PhantomData, ops::Range};
 
 pub struct TchOps<E: tch::kind::Element + Copy + Default> {
     e: PhantomData<E>,
 }
 
 impl<E: tch::kind::Element + Copy + Default> TchOps<E> {
+    pub fn to_device<const D: usize>(
+        tensor: TchTensor<E, D>,
+        device: &LibTorchDevice,
+    ) -> TchTensor<E, D> {
+        let device = (*device).into();
+
+        // We have to manually check if the device is the same, since when it's the case, we need to keep
+        // the same storage reference and not create a new one.
+        if tensor.tensor.device() == device {
+            return tensor;
+        }
+
+        TchTensor::new(tensor.tensor.to(device))
+    }
+
     pub fn reshape<const D1: usize, const D2: usize>(
         tensor: TchTensor<E, D1>,
         shape: Shape<D2>,

burn-tch/src/ops/bool_tensor.rs

@@ -35,7 +35,7 @@ impl<E: TchElement> BoolTensorOps<Self> for LibTorch<E> {
         tensor: TchTensor<bool, D>,
         device: &LibTorchDevice,
     ) -> TchTensor<bool, D> {
-        TchTensor::new(tensor.tensor.to((*device).into()))
+        TchOps::to_device(tensor, device)
     }
 
     fn bool_reshape<const D1: usize, const D2: usize>(

burn-tch/src/ops/int_tensor.rs

@@ -38,7 +38,7 @@ impl<E: TchElement> IntTensorOps<Self> for LibTorch<E> {
         tensor: TchTensor<i64, D>,
         device: &LibTorchDevice,
     ) -> TchTensor<i64, D> {
-        TchTensor::new(tensor.tensor.to((*device).into()))
+        TchOps::to_device(tensor, device)
     }
 
     fn int_reshape<const D1: usize, const D2: usize>(

burn-tch/src/ops/tensor.rs

@@ -102,7 +102,7 @@ impl<E: TchElement> FloatTensorOps<Self> for LibTorch<E> {
         tensor: TchTensor<E, D>,
         device: &LibTorchDevice,
     ) -> TchTensor<E, D> {
-        TchTensor::new(tensor.tensor.to((*device).into()))
+        TchOps::to_device(tensor, device)
     }
 
     fn float_empty<const D: usize>(