Fix candle backend sync (#1579)

* Fix candle backend sync

* tch mps sync

* clippy

---------

Co-authored-by: louisfd <[email protected]>

Author: nathanielsimard

crates/burn-candle/src/backend.rs

@@ -1,6 +1,6 @@
 use std::marker::PhantomData;
 
-use burn_tensor::backend::Backend;
+use burn_tensor::{backend::Backend, Device};
 use candle_core::DeviceLocation;
 
 use crate::{
@@ -91,4 +91,21 @@ impl<F: FloatCandleElement, I: IntCandleElement> Backend for Candle<F, I> {
         // TODO submit an issue at Candle
         panic!("Manual seed not supported by Candle. ")
     }
+
+    fn sync(device: &Device<Self>) {
+        let device: candle_core::Device = (*device).into();
+
+        match device {
+            candle_core::Device::Cpu => (),
+            candle_core::Device::Cuda(device) => {
+                #[cfg(feature = "cuda")]
+                device.synchronize().unwrap();
+            }
+            candle_core::Device::Metal(device) => {
+                // For some reason, device.wait_until_completed() does not seem to work,
+                // and neither does writing and reading a value with into_data
+                panic!("Device synchronization unavailable with Metal device on Candle backend")
+            }
+        }
+    }
 }

crates/burn-tch/src/backend.rs

@@ -3,6 +3,8 @@ use crate::PrecisionBridge;
 use super::element::TchElement;
 use super::TchTensor;
 use burn_tensor::backend::Backend;
+use burn_tensor::ops::IntTensorOps;
+use burn_tensor::{Int, Tensor};
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
 /// The device struct when using the `tch` backend.
@@ -102,10 +104,19 @@ impl<E: TchElement> Backend for LibTorch<E> {
     }
 
     fn sync(device: &Self::Device) {
-        if let LibTorchDevice::Cuda(index) = device {
-            tch::Cuda::synchronize(*index as i64);
-        } else if let LibTorchDevice::Mps = device {
-            panic!("Can't sync MPS device")
+        match device {
+            LibTorchDevice::Cpu => (),
+            LibTorchDevice::Cuda(index) => {
+                tch::Cuda::synchronize(*index as i64);
+            }
+            _ => {
+                // When there is no explicit way to synchronize, we write and read one value to sync
+                Tensor::<Self, 1, Int>::from_primitive(<Self as IntTensorOps<Self>>::int_zeros(
+                    [1].into(),
+                    device,
+                ))
+                .into_data();
+            }
         }
     }
 }