From 7fc0f81836faa6c15aa199959b5da7dbd7bd04ea Mon Sep 17 00:00:00 2001
From: tlk-dsg <467460833@qq.com>
Date: Tue, 28 Sep 2021 20:04:56 +0800
Subject: [PATCH] test

---
 setup.py                                           | 2 +-
 src/deepke/name_entity_re/few_shot/models/model.py | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/setup.py b/setup.py
index 9d65991..af12d55 100644
--- a/setup.py
+++ b/setup.py
@@ -1,7 +1,7 @@
 from setuptools import setup, find_packages
 setup(
     name='deepke',  # 打包后的包文件名
-    version='0.2.60',    #版本号
+    version='0.2.61',    #版本号
     keywords=["pip", "RE","NER","AE"],    # 关键字
     description='DeepKE 是基于 Pytorch 的深度学习中文关系抽取处理套件。',  # 说明
     long_description="client",  #详细说明
diff --git a/src/deepke/name_entity_re/few_shot/models/model.py b/src/deepke/name_entity_re/few_shot/models/model.py
index d7d2d9f..2f8b41f 100644
--- a/src/deepke/name_entity_re/few_shot/models/model.py
+++ b/src/deepke/name_entity_re/few_shot/models/model.py
@@ -407,7 +407,7 @@ def _no_beam_search_generate(decoder: PromptBartDecoder, state, tokens=None, max
                              eos_token_id=None,
                              repetition_penalty=1.0, length_penalty=1.0, pad_token_id=0,
                              restricter=None):
-    device = _get_model_device(decoder)
+    device = get_model_device(decoder)
     if tokens is None:
         if bos_token_id is None:
             raise RuntimeError("You have to specify either `tokens` or `bos_token_id`.")
@@ -492,7 +492,7 @@ def _beam_search_generate(decoder: PromptBartDecoder, tokens=None, state=None, m
                           restricter=None) -> torch.LongTensor:
     assert do_sample is False
     # beam search
-    device = _get_model_device(decoder)
+    device = get_model_device(decoder)
     if tokens is None:
         if bos_token_id is None:
             raise RuntimeError("You have to specify either `tokens` or `bos_token_id`.")