mend

mims-harvard · Jan 21, 2025 · 90c98aa · 90c98aa
1 parent ba7602a
commit 90c98aa
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/tdc/test/test_model_server.py b/tdc/test/test_model_server.py
@@ -81,6 +81,8 @@ def testGeneformerPerturb(self):
                     cell = np.append(cell, 0)
                 batch[idx] = cell
         input_tensor = torch.tensor(batch)
+        assert input_tensor.shape[0] == 512, "unexpected batch size"
+        assert input_tensor.shape[1] == mdim, f"unexpected gene length {mdim}"
         attention_mask = torch.tensor([[t != 0 for t in cell] for cell in batch
                                       ])
         outputs = model(input_tensor,