fix: move __cast_dtype__

jina-ai · Dec 8, 2022 · edf4629 · edf4629
1 parent fc36e9b
commit edf4629
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 4 deletions.
diff --git a/server/clip_server/executors/clip_torch.py b/server/clip_server/executors/clip_torch.py
@@ -12,14 +12,16 @@
     set_rank,
     split_img_txt_da,
 )
-from clip_server.helper import __cast_dtype__
 from clip_server.model import clip
 from clip_server.model.clip_model import CLIPModel
 from clip_server.model.tokenization import Tokenizer
 from jina import DocumentArray, Executor, requests
 from opentelemetry.trace import NoOpTracer, Span
 
 
+__cast_dtype__ = {'fp16': torch.float16, 'fp32': torch.float32, 'bf16': torch.bfloat16}
+
+
 class CLIPEncoder(Executor):
     def __init__(
         self,

diff --git a/server/clip_server/helper.py b/server/clip_server/helper.py
@@ -20,9 +20,6 @@
 )
 
 
-__cast_dtype__ = {'fp16': torch.float16, 'fp32': torch.float32, 'bf16': torch.bfloat16}
-
-
 def _version_check(package: str = None, github_repo: str = None):
     try:
-Original file line number
+Diff line change
@@ Expand Up / @@ -20,9 +20,6 @@ @@
     )
-    __cast_dtype__ = {'fp16': torch.float16, 'fp32': torch.float32, 'bf16': torch.bfloat16}
     def _version_check(package: str = None, github_repo: str = None):
         try:
@@ Expand Down @@