diff --git a/libnd4j/tests_cpu/layers_tests/DataBufferTests.cpp b/libnd4j/tests_cpu/layers_tests/DataBufferTests.cpp
index e6d270141..03d7bb38e 100644
--- a/libnd4j/tests_cpu/layers_tests/DataBufferTests.cpp
+++ b/libnd4j/tests_cpu/layers_tests/DataBufferTests.cpp
@@ -48,7 +48,7 @@ TEST_F(DataBufferTests, test_alloc_limit_1) {
     auto odUse = MemoryCounter::getInstance()->allocatedDevice(deviceId);
     auto ogUse = MemoryCounter::getInstance()->allocatedGroup(MemoryType::HOST);
 
-    auto limitSize = 150 * 1024 * 1024;
+    auto limitSize = odUse + (150 * 1024 * 1024);
     auto allocSize = 100000000;
 
     MemoryCounter::getInstance()->setDeviceLimit(deviceId, odLimit + limitSize);
diff --git a/libnd4j/tests_cpu/layers_tests/DataBufferTestsCuda.cu b/libnd4j/tests_cpu/layers_tests/DataBufferTestsCuda.cu
index 355633a4e..9ac68d55e 100644
--- a/libnd4j/tests_cpu/layers_tests/DataBufferTestsCuda.cu
+++ b/libnd4j/tests_cpu/layers_tests/DataBufferTestsCuda.cu
@@ -51,7 +51,7 @@ TEST_F(DataBufferTestsCuda, test_alloc_limit_1) {
     auto opUse = MemoryCounter::getInstance()->allocatedGroup(MemoryType::HOST);
     auto osUse = MemoryCounter::getInstance()->allocatedGroup(MemoryType::DEVICE);
 
-    auto limitSize = 150000000;
+    auto limitSize = odUse + 150000000;
     auto allocSize = 100000000;
 
     MemoryCounter::getInstance()->setDeviceLimit(deviceId, odLimit + limitSize);
@@ -72,7 +72,7 @@ TEST_F(DataBufferTestsCuda, test_alloc_limit_1) {
 
     // this allocation should fail, since we're allocating too much
     try {
-        DataBuffer bufferFailed(allocSize, DataType::INT32);
+        DataBuffer bufferFailed(allocSize + 1, DataType::INT32);
         ASSERT_TRUE(false);
     } catch (allocation_exception &e) {
         // we expect exception here