Add model initialization tests

atong01 · Nov 23, 2023 · 1b74617 · 1b74617
1 parent 91c5dfd
commit 1b74617
Showing 1 changed file with 14 additions and 0 deletions.
diff --git a/tests/test_models.py b/tests/test_models.py
@@ -0,0 +1,14 @@
+import pytest
+from torchcfm.models import MLP
+from torchcfm.models.unet import UNetModel
+
+
+def test_initialize_models():
+    model = UNetModel(
+        dim=(1, 28, 28),
+        num_channels=32,
+        num_res_blocks=1,
+        num_classes=10,
+        class_cond=True,
+    )
+    model = MLP(dim=2, time_varying=True, w=64)