🐞 Fix GMM test (open-edge-platform#1696)

* Separate data

Signed-off-by: Ashwin Vaidya <[email protected]>

* limit torch version

Signed-off-by: Ashwin Vaidya <[email protected]>

---------

Signed-off-by: Ashwin Vaidya <[email protected]>

Author: ashwinvaidya17

requirements/base.txt

@@ -12,6 +12,7 @@ pandas>=1.1.0
 lightning
 setuptools>=41.0.0
 timm>=0.5.4,<=0.6.13
+torch>=2,<2.2.0 # rkde export fails even with ONNX 17 (latest) with torch 2.2.0. TODO(ashwinvaidya17): revisit
 torchmetrics==0.10.3
 rich-argparse
 open-clip-torch>=2.23.0

tests/unit/models/components/clustering/test_gmm.py

@@ -16,8 +16,8 @@ def test_fit_and_predict() -> None:
     # Create some synthetic data
     data = torch.cat(
         [
-            torch.randn(100, 2) + torch.tensor([2.0, 2.0]),
-            torch.randn(100, 2) + torch.tensor([-2.0, -2.0]),
+            torch.randn(100, 2) + torch.tensor([10.0, 10.0]),
+            torch.randn(100, 2) + torch.tensor([-10.0, -10.0]),
         ],
     )