未验证 提交 353a4c10 编写于 作者: M Michael Wyatt 提交者: GitHub

Update tests/unit/inference/test_inference.py

上级 8552bd6a
...@@ -391,12 +391,12 @@ class TestLowCpuMemUsage(DistributedTest): ...@@ -391,12 +391,12 @@ class TestLowCpuMemUsage(DistributedTest):
def test( def test(
self, self,
model_w_task, model_w_task,
dtype,
query, query,
inf_kwargs, inf_kwargs,
assert_fn, assert_fn,
): ):
model, task = model_w_task model, task = model_w_task
dtype = torch.float16
local_rank = int(os.getenv("LOCAL_RANK", "0")) 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") pipe = pipeline(task, model=model, model_kwargs={"low_cpu_mem_usage": True}, device=local_rank, framework="pt")
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册