From 353a4c10fde0ff21d2f78e32852425f2059dd1dc Mon Sep 17 00:00:00 2001
From: Michael Wyatt <michaelwyatt@microsoft.com>
Date: Tue, 22 Aug 2023 10:49:05 -0700
Subject: [PATCH] Update tests/unit/inference/test_inference.py

---
 tests/unit/inference/test_inference.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/unit/inference/test_inference.py b/tests/unit/inference/test_inference.py
index 63f67cbc..da5ad81e 100644
--- a/tests/unit/inference/test_inference.py
+++ b/tests/unit/inference/test_inference.py
@@ -391,12 +391,12 @@ class TestLowCpuMemUsage(DistributedTest):
     def test(
         self,
         model_w_task,
-        dtype,
         query,
         inf_kwargs,
         assert_fn,
     ):
         model, task = model_w_task
+        dtype = torch.float16
         local_rank = int(os.getenv("LOCAL_RANK", "0"))
 
         pipe = pipeline(task, model=model, model_kwargs={"low_cpu_mem_usage": True}, device=local_rank, framework="pt")
-- 
GitLab