correct assert

Author: kevbuh

.github/workflows/ci.yml

@@ -2,9 +2,9 @@ name: CI
 
 on:
   push:
-    branches: [ main ]
+    branches: [main]
   pull_request:
-    branches: [ main ]
+    branches: [main]
 
 jobs:
   test:
@@ -33,7 +33,7 @@ jobs:
           if [ -f requirements.txt ]; then
             pip install -r requirements.txt
           else
-            pip install torch pytest
+            pip install torch pytest numpy
           fi
 
       - name: Run tests

bitnet/model.py

@@ -144,7 +144,7 @@ def forward(self, x):
 
 class BitNet(nn.Module):
   def __init__(self, vocab_size, d_model, block_size, n_layer, n_head, n_kv_head, ffn_dim):
-    assert(n_head % n_kv_head == 0, "n_head must be divisible by n_kv_head")
+    if n_head % n_kv_head != 0: raise ValueError("n_head must be divisible by n_kv_head")
     super().__init__()
     self.block_size = block_size
     self.embed_tokens = nn.Embedding(vocab_size, d_model)