reshape fix

tdc/test/test_model_server.py

@@ -104,10 +104,17 @@ def testGeneformerTokenizer(self):
         model = geneformer.load()
         input_tensor = torch.tensor(cells)
         input_tensor = torch.squeeze(input_tensor)
-        raise Exception("shape is", input_tensor.shape, "values are\n", input_tensor)
-        out = model(input_tensor)
+        x = input_tensor.shape[0]
+        y = input_tensor.shape[1]
+        input_tensor = input_tensor.reshape(x, y)
+        out = None  # try-except block
+        try:
+            out = model(input_tensor)
+        except Exception as e:
+            raise Exception("shape is", input_tensor.shape, "exception was: {}".format(e), "values are\n", input_tensor)
         assert out, "FAILURE: Geneformer output is false-like. Value = {}".format(out)
-        assert len(out) == len(cells), "FAILURE: Geneformer output and cells input don't have the same length. {} vs {}".format(len(out), len(cells))
+        assert out.shape[0] == input_tensor.shape[0], "FAILURE: Geneformer output and input tensor input don't have the same length. {} vs {}".format(out.shape[0], input_tensor.shape[0])
+        assert out.shape[0] == len(cells), "FAILURE: Geneformer output and tokenized cells don't have the same length. {} vs {}".format(out.shape[0], len(cells))
 
     def tearDown(self):
         try: