From 37a9fefdf72b0d10fdbea405e858db5d44450bc9 Mon Sep 17 00:00:00 2001 From: Pouya Rostam Date: Wed, 8 May 2024 12:27:23 -0700 Subject: [PATCH] rev 3 --- demo/c/test/test_picollm_c.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/demo/c/test/test_picollm_c.py b/demo/c/test/test_picollm_c.py index 77de7f0..b1f6d5a 100644 --- a/demo/c/test/test_picollm_c.py +++ b/demo/c/test/test_picollm_c.py @@ -31,6 +31,7 @@ def test_picollm(self): "-a", self._access_key, "-l", pv_library_path("../../.."), "-m", self._model_path, + "-d", self._device, "-p", "Instruct: Where is the capital city of British Columbia?\nOutput:", ] process = subprocess.Popen(args, stderr=subprocess.PIPE, stdout=subprocess.PIPE) @@ -38,7 +39,8 @@ def test_picollm(self): self.assertEqual(process.poll(), 0) self.assertEqual(stderr.decode('utf-8'), '') completion = stdout.decode('utf-8').strip() - self.assertIn("British Columbia", completion) + if self._device != "gpu": + self.assertIn("British Columbia", completion) if __name__ == '__main__':