Update README.md

README.md

@@ -1130,7 +1130,7 @@ def mean_pooling(model_output, attention_mask):
 sentences = ['How is the weather today?', '今天天气怎么样?']
 
 tokenizer = AutoTokenizer.from_pretrained('jinaai/jina-embeddings-v2-base-zh')
-model = AutoModel.from_pretrained('jinaai/jina-embeddings-v2-base-zh', trust_remote_code=True)
+model = AutoModel.from_pretrained('jinaai/jina-embeddings-v2-base-zh', trust_remote_code=True, torch_dtype=torch.bfloat16)
 
 encoded_input = tokenizer(sentences, padding=True, truncation=True, return_tensors='pt')
 
@@ -1148,11 +1148,12 @@ You can use Jina Embedding models directly from transformers package.
 
 ```python
 !pip install transformers
+import torch
 from transformers import AutoModel
 from numpy.linalg import norm
 
 cos_sim = lambda a,b: (a @ b.T) / (norm(a)*norm(b))
-model = AutoModel.from_pretrained('jinaai/jina-embeddings-v2-base-zh', trust_remote_code=True) # trust_remote_code is needed to use the encode method
+model = AutoModel.from_pretrained('jinaai/jina-embeddings-v2-base-zh', trust_remote_code=True, torch_dtype=torch.bfloat16)
 embeddings = model.encode(['How is the weather today?', '今天天气怎么样?'])
 print(cos_sim(embeddings[0], embeddings[1]))
 ```