diff --git a/tests/kernels/test_layernorm.py b/tests/kernels/test_layernorm.py
index 4f65a14b2f..0621b5d9ea 100644
--- a/tests/kernels/test_layernorm.py
+++ b/tests/kernels/test_layernorm.py
@@ -3,7 +3,7 @@ import torch
 import triton
 
 from unsloth.kernels.layernorm import LayerNorm
-from tests.conftest import set_seed 
+from tests.conftest import set_seed, assert_all_close
 
 # Fixture for test matrices and associated parameters
 @set_seed()
@@ -30,9 +30,8 @@ def test_layer_norm_forward(test_data):
     pytorch_layer_norm.bias = torch.nn.Parameter(bias)
     pytorch_output = pytorch_layer_norm(x)
 
-    # Check if outputs are close
-    assert torch.allclose(triton_output, pytorch_output, rtol=1e-05, atol=1e-08), \
-           "Forward pass outputs differ between Triton and PyTorch."
+    # Check if outputs are close using assert_all_close
+    assert_all_close(triton_output, pytorch_output, rtol=1e-05, atol=1e-08)
 
 
 def test_layer_norm_backward(test_data):
@@ -51,6 +50,5 @@ def test_layer_norm_backward(test_data):
     pytorch_output.sum().backward()
     pytorch_grad = x.grad
 
-    # Check if gradients are close
-    assert torch.allclose(triton_grad, pytorch_grad, rtol=1e-05, atol=1e-08), \
-           "Backward pass gradients differ between Triton and PyTorch."
+    # Check if gradients are close using assert_all_close
+    assert_all_close(triton_grad, pytorch_grad, rtol=1e-05, atol=1e-08)
diff --git a/tests/kernels/test_relu.py b/tests/kernels/test_relu.py
index 3c2a09ec34..3785deab3a 100644
--- a/tests/kernels/test_relu.py
+++ b/tests/kernels/test_relu.py
@@ -2,8 +2,8 @@ import pytest
 import torch
 import triton
 
-from unsloth.kernels.relu import relu_kernel  # Import your relu_kernel function
-from tests.conftest import set_seed
+from unsloth.kernels.relu import relu_kernel
+from tests.conftest import set_seed, assert_all_close
 
 @set_seed
 @pytest.fixture(params=[(100, 100), (1024, 1024), (5000, 1024), (12345, 5678)])
@@ -21,7 +21,5 @@ def test_relu_kernel(test_matrix):
     torch_relu = torch.nn.ReLU()
     torch_output = torch_relu(test_matrix)
 
-    # Check if the outputs are close enough
-    # You can adjust rtol and atol based on the precision you expect
-    assert torch.allclose(triton_output, torch_output, rtol=1e-05, atol=1e-08), \
-           "The outputs are not close enough between Triton and PyTorch implementation."
+    # Check if the outputs are close enough using assert_all_close
+    assert_all_close(triton_output, torch_output, rtol=1e-05, atol=1e-08)