fix: check dtype when loading models (#872)

* fix: check dtype when loading models * fix: black
jina-ai · Dec 7, 2022 · fd16e5a · fd16e5a
1 parent c7af9f7
commit fd16e5a
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/server/clip_server/model/model.py b/server/clip_server/model/model.py
@@ -464,7 +464,9 @@ def load_openai_model(
 
         # model from OpenAI state dict is in manually cast fp16 mode, must be converted for AMP/fp32/bf16 use
         model = model.to(device)
-        if dtype == torch.float32 or dtype.startswith('amp'):
+        if dtype == torch.float32 or (
+            isinstance(dtype, str) and dtype.startswith('amp')
+        ):
             model.float()
         elif dtype == torch.bfloat16:
             convert_weights_to_lp(model, dtype=torch.bfloat16)