added comments to main test file

jshuadvd · Jun 12, 2024 · 985e6e6 · 985e6e6
1 parent 232fa37
commit 985e6e6
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/tests/test_main.py b/tests/test_main.py
@@ -50,6 +50,7 @@ def test_longrope_model_initialization():
     assert model.max_len == 65536
 
 
+# Testing the LongRoPEModel class with embedding
 def test_longrope_model_embedding():
     model = LongRoPEModel(
         d_model=512, n_heads=8, num_layers=6, vocab_size=50257, max_len=65536
@@ -60,6 +61,7 @@ def test_longrope_model_embedding():
     assert not torch.equal(input_ids, embeddings)
 
 
+# Testing the LongRoPEModel class with transformers
 def test_longrope_model_transformers():
     model = LongRoPEModel(
         d_model=512, n_heads=8, num_layers=6, vocab_size=50257, max_len=65536
@@ -71,6 +73,7 @@ def test_longrope_model_transformers():
     assert embeddings.shape == (2, 1024, 512)
 
 
+# Testing the LongRoPEModel class with forward
 def test_longrope_model_forward():
     model = LongRoPEModel(
         d_model=512, n_heads=8, num_layers=6, vocab_size=50257, max_len=65536