From d8984045006e2e8e8bbafb40bbbd693b9a7cd564 Mon Sep 17 00:00:00 2001 From: Lev Kurilenko Date: Thu, 1 Dec 2022 19:39:09 +0000 Subject: [PATCH] Set replace_with_kernel_inject=True in BERT fill-mask example --- inference/huggingface/fill-mask/test-bert.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/inference/huggingface/fill-mask/test-bert.py b/inference/huggingface/fill-mask/test-bert.py index 994d5fde2..441c617b7 100644 --- a/inference/huggingface/fill-mask/test-bert.py +++ b/inference/huggingface/fill-mask/test-bert.py @@ -14,7 +14,7 @@ pipe.model, mp_size=world_size, dtype=torch.float, - injection_policy={BertLayer : ('output.dense')} + replace_with_kernel_inject=True ) pipe.device = torch.device(f'cuda:{local_rank}')