Prepare tests for memory allocation while estimating

Author: mfep

tests/test_prep/test_stripe.py

@@ -96,14 +96,15 @@ def test_remove_stripe_fw_calc_mem(slices, level, dim_x, wname, ensure_clean_mem
         remove_stripe_fw(cp.copy(data), wname=wname, level=level)
     actual_mem_peak = hook.max_mem
 
-    hook = MaxMemoryHook()
-    with hook:
+    try:
         estimated_mem_peak = remove_stripe_fw(
             data.shape, level=level, wname=wname, calc_peak_gpu_mem=True
         )
+    except cp.cuda.memory.OutOfMemoryError:
+        pytest.skip("Not enough GPU memory to estimate memory peak")
 
     assert actual_mem_peak * 0.99 <= estimated_mem_peak
-    assert estimated_mem_peak <= actual_mem_peak * 1.3
+    assert estimated_mem_peak <= actual_mem_peak * 1.15
 
 
 @pytest.mark.parametrize("wname", ["haar", "db4", "sym5", "sym16", "bior4.4"])
@@ -115,11 +116,12 @@ def test_remove_stripe_fw_calc_mem_big(wname, slices, level, ensure_clean_memory
     dim_y = 901
     dim_x = 1200
     data_shape = (slices, dim_x, dim_y)
-    hook = MaxMemoryHook()
-    with hook:
+    try:
         estimated_mem_peak = remove_stripe_fw(
             data_shape, wname=wname, level=level, calc_peak_gpu_mem=True
         )
+    except cp.cuda.memory.OutOfMemoryError:
+        pytest.skip("Not enough GPU memory to estimate memory peak")
     av_mem = cp.cuda.Device().mem_info[0]
     if av_mem < estimated_mem_peak:
         pytest.skip("Not enough GPU memory to run this test")
@@ -131,7 +133,7 @@ def test_remove_stripe_fw_calc_mem_big(wname, slices, level, ensure_clean_memory
     actual_mem_peak = hook.max_mem
 
     assert actual_mem_peak * 0.99 <= estimated_mem_peak
-    assert estimated_mem_peak <= actual_mem_peak * 1.3
+    assert estimated_mem_peak <= actual_mem_peak * 1.15
 
 
 @pytest.mark.parametrize("angles", [180, 181])