diff --git a/setup.py b/setup.py
index 7de4287e05c..113c1c041cd 100644
--- a/setup.py
+++ b/setup.py
@@ -57,6 +57,7 @@
     "sphinx-markdown-tables>=0.0.15",
     "wheel>=0.36.2",
     "pytest>=6.0.0",
+    "flaky>=3.0.0",
     "sphinx-rtd-theme",
 ]
 
diff --git a/tests/sparseml/tensorflow_v1/optim/test_analyzer_module.py b/tests/sparseml/tensorflow_v1/optim/test_analyzer_module.py
index 8b51e347af7..ebebe0f05c5 100644
--- a/tests/sparseml/tensorflow_v1/optim/test_analyzer_module.py
+++ b/tests/sparseml/tensorflow_v1/optim/test_analyzer_module.py
@@ -85,6 +85,7 @@ def resnet_v2_50(init_weights):
         return tf_compat.get_default_graph()
 
 
+@pytest.mark.flaky
 @pytest.mark.skipif(
     os.getenv("NM_ML_SKIP_TENSORFLOW_TESTS", False),
     reason="Skipping tensorflow_v1 tests",
diff --git a/tests/sparseml/tensorflow_v1/optim/test_mask_creator_pruning.py b/tests/sparseml/tensorflow_v1/optim/test_mask_creator_pruning.py
index 92aa5a6128d..082c90f4578 100644
--- a/tests/sparseml/tensorflow_v1/optim/test_mask_creator_pruning.py
+++ b/tests/sparseml/tensorflow_v1/optim/test_mask_creator_pruning.py
@@ -26,6 +26,7 @@
 from sparseml.tensorflow_v1.utils import eval_tensor_sparsity, tf_compat
 
 
+@pytest.mark.flaky
 @pytest.mark.parametrize(
     ("tensor_shape,mask_creator"),
     [
diff --git a/tests/sparseml/tensorflow_v1/optim/test_mask_pruning.py b/tests/sparseml/tensorflow_v1/optim/test_mask_pruning.py
index 31a2019cb15..ecbda8a2bfc 100644
--- a/tests/sparseml/tensorflow_v1/optim/test_mask_pruning.py
+++ b/tests/sparseml/tensorflow_v1/optim/test_mask_pruning.py
@@ -202,6 +202,7 @@ def test_create_op_pruning_conv(sparsity_val: float, mask_creator: PruningMaskCr
                 assert sess.run(mask_vals_are_grouped)
 
 
+@pytest.mark.flaky
 @pytest.mark.skipif(
     os.getenv("NM_ML_SKIP_TENSORFLOW_TESTS", False),
     reason="Skipping tensorflow_v1 tests",
@@ -391,6 +392,7 @@ def test_apply_op_vars_masks(
                 assert abs(var_sparsity - sparsity_val) < 1e-2
 
 
+@pytest.mark.flaky
 @pytest.mark.skipif(
     os.getenv("NM_ML_SKIP_TENSORFLOW_TESTS", False),
     reason="Skipping tensorflow_v1 tests",
@@ -502,6 +504,7 @@ def _expected_sparsity(
     )
 
 
+@pytest.mark.flaky
 @pytest.mark.skipif(
     os.getenv("NM_ML_SKIP_TENSORFLOW_TESTS", False),
     reason="Skipping tensorflow_v1 tests",