cleanup

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>

Author: kylesayrs

tests/test_transform/factory/test_correctness.py

@@ -23,12 +23,13 @@
 from tests.testing_utils import requires_accelerate, requires_gpu
 
 
-_test_schemes = [
-    TransformScheme(type=name) for name in TransformFactory.registered_names()
-] + [
-    TransformScheme(type=name, randomize=True)
-    for name in TransformFactory.registered_names()
-]
+def all_schemes():
+    base = [TransformScheme(type=name) for name in TransformFactory.registered_names()]
+    randomized = [
+        TransformScheme(type=name, randomize=True)
+        for name in TransformFactory.registered_names()
+    ]
+    return base + randomized
 
 
 class TransformableModel(torch.nn.Module):
@@ -45,7 +46,7 @@ def forward(self, x):
         return x
 
 
-@pytest.mark.parametrize("scheme", _test_schemes)
+@pytest.mark.parametrize("scheme", all_schemes())
 def test_correctness_linear(scheme):
     size = (4, 8)
     module = torch.nn.Linear(*size, bias=True)
@@ -72,7 +73,7 @@ def test_correctness_linear(scheme):
     assert torch.allclose(true_output, output, atol=1e-5, rtol=0.0)
 
 
-@pytest.mark.parametrize("scheme", _test_schemes)
+@pytest.mark.parametrize("scheme", all_schemes())
 def test_correctness_model(scheme, offload=False):
     # load model
     model = TransformableModel(2, 4, 8, 16, 32, 64)
@@ -110,6 +111,6 @@ def test_correctness_model(scheme, offload=False):
 
 @requires_gpu
 @requires_accelerate()
-@pytest.mark.parametrize("scheme", _test_schemes)
+@pytest.mark.parametrize("scheme", all_schemes())
 def test_correctness_model_offload(scheme):
     test_correctness_model(scheme, offload=True)