cleanup 2

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

Author: kylesayrs

tests/test_transform/factory/test_memory.py

@@ -26,12 +26,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):
@@ -48,7 +49,7 @@ def forward(self, x):
         return x
 
 
-@pytest.mark.parametrize("scheme", _test_schemes)
+@pytest.mark.parametrize("scheme", all_schemes())
 def test_memory_sharing(scheme, offload=False):
     # load scheme and factory
     scheme = TransformScheme(
@@ -98,12 +99,12 @@ def test_memory_sharing(scheme, offload=False):
 
 @requires_gpu
 @requires_accelerate()
-@pytest.mark.parametrize("scheme", _test_schemes)
+@pytest.mark.parametrize("scheme", all_schemes())
 def test_memory_sharing_offload(scheme):
     test_memory_sharing(scheme, offload=True)
 
 
-@pytest.mark.parametrize("scheme", _test_schemes)
+@pytest.mark.parametrize("scheme", all_schemes())
 def test_memory_sharing_training(scheme):
     scheme.requires_grad = True
     test_memory_sharing(scheme, offload=False)