chore: adding conditioner tests

ml-gde · Oct 7, 2024 · 4a7461a · 4a7461a
1 parent 55abdee
commit 4a7461a
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 1 deletion.
diff --git a/jflux/modules/conditioner.py b/jflux/modules/conditioner.py
@@ -2,9 +2,9 @@
 from flax import nnx
 from transformers import (
     CLIPTokenizer,
+    T5Tokenizer,
     FlaxCLIPTextModel,
     FlaxT5EncoderModel,
-    T5Tokenizer,
 )
 
 

diff --git a/tests/modules/test_conditioner.py b/tests/modules/test_conditioner.py
@@ -0,0 +1,10 @@
+import numpy as np
+
+from flux.modules.conditioner import HFEmbedder as TorchHFEmbedder
+from jflux.modules.conditioner import HFEmbedder as JaxHFEmbedder
+
+
+class HFEmbedderTestCase(np.testing.TestCase):
+    def test_hf_embed(self):
+        # initialize layers
+        TorchHFEmbedder()