Fix cast float test precision (tracel-ai#2586)

tsanona · Dec 9, 2024 · 19975e9 · 19975e9
1 parent 170b3bb
commit 19975e9
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/crates/burn-tensor/src/tests/ops/cast.rs b/crates/burn-tensor/src/tests/ops/cast.rs
@@ -46,6 +46,7 @@ mod tests {
         let output = tensor.cast(DType::F32);
 
         assert_eq!(output.dtype(), DType::F32);
-        output.into_data().assert_approx_eq(&data, 5);
+        // Use precision 2 for parametrized tests in f16 and bf16
+        output.into_data().assert_approx_eq(&data, 2);
     }
 }