add alloc extern crate

gonzalesMK · Sep 20, 2023 · 8619aa3 · 8619aa3
1 parent bb84984
commit 8619aa3
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 0 deletions.
diff --git a/burn-candle/src/lib.rs b/burn-candle/src/lib.rs
@@ -15,6 +15,7 @@ pub use tensor::*;
 
 #[cfg(test)]
 mod tests {
+    extern crate alloc;
     use super::*;
 
     pub type TestBackend = CandleBackend<f32, i64>;

diff --git a/burn-tch/src/lib.rs b/burn-tch/src/lib.rs
@@ -14,6 +14,8 @@ pub use tensor::*;
 
 #[cfg(test)]
 mod tests {
+    extern crate alloc;
+
     type TestBackend = crate::TchBackend<f32>;
     type TestTensor<const D: usize> = burn_tensor::Tensor<TestBackend, D>;
     type TestTensorInt<const D: usize> = burn_tensor::Tensor<TestBackend, D, burn_tensor::Int>;

diff --git a/burn-tensor/src/tests/ops/cat.rs b/burn-tensor/src/tests/ops/cat.rs
@@ -1,6 +1,7 @@
 #[burn_tensor_testgen::testgen(cat)]
 mod tests {
     use super::*;
+    use alloc::vec::Vec;
     use burn_tensor::{Bool, Data, Int, Tensor};
 
     #[test]
@@ -57,4 +58,12 @@ mod tests {
         let data_expected = Data::from([[[1.0, 2.0, 3.0]], [[1.1, 2.1, 3.1]], [[4.0, 5.0, 6.0]]]);
         data_expected.assert_approx_eq(&data_actual, 3);
     }
+
+    #[test]
+    #[should_panic]
+    fn should_panic_when_list_of_vectors_is_empty() {
+        let tensor: Vec<TestTensor<2>> = vec![];
+
+        TestTensor::cat(tensor, 0).into_data();
+    }
 }
diff --git a/burn-wgpu/src/lib.rs b/burn-wgpu/src/lib.rs
@@ -4,6 +4,7 @@
 
 #[macro_use]
 extern crate derive_new;
+extern crate alloc;
 
 mod ops;