Update ARAD constructor test, change default to v2.

sandialabs · Jan 15, 2024 · 7d33209 · 7d33209
1 parent 03c06ba
commit 7d33209
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 5 deletions.
diff --git a/riid/models/neural_nets/arad.py b/riid/models/neural_nets/arad.py
@@ -238,7 +238,7 @@ def call(self, x):
 class ARAD(PyRIIDModel):
     """PyRIID-compatible ARAD model to work with SampleSets.
     """
-    def __init__(self, arad_version: str = "v1", latent_dim=None):
+    def __init__(self, arad_version: str = "v2", latent_dim=None):
         """
         Args:
             arad_version: version of ARAD to use, can be "v1" or "v2".

diff --git a/tests/model_tests.py b/tests/model_tests.py
@@ -88,10 +88,8 @@ def test_all_constructors(self):
         _ = MLPClassifier()
         _ = LabelProportionEstimator()
         _ = MultiEventClassifier()
-        arad_v1 = ARADv1TF()
-        _ = ARAD(arad_v1)
-        arad_v2 = ARADv2TF()
-        _ = ARAD(arad_v2)
+        _ = ARAD(arad_version="v1")
+        _ = ARAD(arad_version="v2")
 
 
 if __name__ == "__main__":