squeeze last dim in test

mims-harvard · Oct 26, 2024 · 741555f · 741555f
1 parent b913742
commit 741555f
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/tdc/test/test_model_server.py b/tdc/test/test_model_server.py
@@ -106,6 +106,7 @@ def testGeneformerTokenizer(self):
         input_tensor = torch.tensor(cells)
         input_tensor = torch.squeeze(input_tensor)
         try:
+            input_tensor.squeeze(2)  # last dim is zero
             out = model(input_tensor)
         except Exception as e:
             raise Exception("tensor shape is", input_tensor.shape, "exception was:", e)