xfail quantization test

Lightning-AI · May 30, 2022 · 31a4b6d · 31a4b6d
1 parent daaff61
commit 31a4b6d
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/tests/callbacks/test_quantization.py b/tests/callbacks/test_quantization.py
@@ -22,6 +22,7 @@
 from pytorch_lightning import seed_everything, Trainer
 from pytorch_lightning.callbacks import QuantizationAwareTraining
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
+from pytorch_lightning.utilities.imports import _TORCH_GREATER_EQUAL_1_11
 from pytorch_lightning.utilities.memory import get_model_size_mb
 from tests.helpers.boring_model import RandomDataset
 from tests.helpers.datamodules import RegressDataModule
@@ -35,6 +36,8 @@
 @RunIf(quantization=True)
 def test_quantization(tmpdir, observe: str, fuse: bool, convert: bool):
     """Parity test for quant model."""
+    pytest.mark.xfail(observe == "average", not fuse, _TORCH_GREATER_EQUAL_1_11, reason="TODO: flakiness in GPU CI")
+
     seed_everything(42)
     dm = RegressDataModule()
     accelerator = "gpu" if torch.cuda.is_available() else "cpu"