update workflow

Author: Daniel Mann

.github/workflows/unit_tests.yml

@@ -9,7 +9,9 @@ on:
       - main
 jobs:
   test-jobs:
-    runs-on: ubuntu-latest
+    runs-on:
+      group: tj-runner-group
+      labels: tj-runners
     steps:
       - uses: actions/checkout@v4
         with:
@@ -23,4 +25,4 @@ jobs:
           pip install -r requirements.txt
       - name: Test Ops
         run: |
-          python -m unittest -v test
+          python -m unittest -s discover -s test -v

test/test_fast_viterbi.py

@@ -14,7 +14,7 @@ def test_best_sequence(self):
                 [0.1, 0.9],
                 [0.1, 0.9],
             ], device="cuda", dtype=torch.float32)
-            .unsqueeze(1)
+            .unsqueeze(0)
             .log()
         )
         edges = (
@@ -33,8 +33,7 @@ def test_best_sequence(self):
         fsa = (2, edges, weights, start_end_states)
 
         output, scores = align_viterbi(log_probs, fsa, seq_lens)
-        best_sequence = list(output[:,0])
-        score = float(scores[0])
+        best_sequence = list(output[0,:])
 
         self.assertEqual(best_sequence, [0, 0, 1, 1, 1])
 

test/test_fbw2.py

@@ -1,5 +1,4 @@
 import itertools
-import gc
 import unittest
 import torch