Update app.py

app.py

@@ -14,7 +14,12 @@ You can also use 🐣e5-mistral🛌🏻 by cloning this space. 🧬🔬🔍 Simp
 Join us : 🌟TeamTonic🌟 is always making cool demos! Join our active builder's🛠️community on 👻Discord:  [![Let's build the future of AI together! 🚀🤖](https://discordapp.com/api/guilds/1109943800132010065/widget.png)](https://discord.gg/GWpVpekp) On 🤗Huggingface: [TeamTonic](https://huggingface.co/TeamTonic) & [MultiTransformer](https://huggingface.co/MultiTransformer) On 🌐Github: [Polytonic](https://github.com/tonic-ai) & contribute to 🌟 [Poly](https://github.com/tonic-ai/poly)
 """
 
-os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:50'
+# os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:50'
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# torch.backends.cuda.matmul.allow_tf32 = True
+# torch.backends.cudnn.allow_tf32 = True
+# torch.backends.cudnn.benchmark = True
 
 def last_token_pool(last_hidden_states: Tensor, attention_mask: Tensor) -> Tensor:
     left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
@@ -31,20 +36,15 @@ def get_detailed_instruct(task_description: str, query: str) -> str:
 
 @spaces.GPU
 def compute_embeddings(*input_texts):
-    torch.backends.cuda.matmul.allow_tf32 = True
-    torch.backends.cudnn.allow_tf32 = True
-    torch.backends.cudnn.benchmark = True
     tokenizer = AutoTokenizer.from_pretrained('intfloat/e5-mistral-7b-instruct')
-    model = AutoModel.from_pretrained('intfloat/e5-mistral-7b-instruct')
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model.to(device)
+    model = AutoModel.from_pretrained('intfloat/e5-mistral-7b-instruct', torch_dtype="auto", device_map=device))
     max_length = 4096
     task = 'Given a web search query, retrieve relevant passages that answer the query'
-
     processed_texts = [get_detailed_instruct(task, text) for text in input_texts]
     batch_dict = tokenizer(processed_texts, max_length=max_length - 1, return_attention_mask=False, padding=False, truncation=True)
     batch_dict['input_ids'] = [input_ids + [tokenizer.eos_token_id] for input_ids in batch_dict['input_ids']]
     batch_dict = tokenizer.pad(batch_dict, padding=True, return_attention_mask=True, return_tensors='pt')
+    batch_dict = {k: v.to(device) for k, v in batch_dict.items()}
     outputs = model(**batch_dict)
     embeddings = last_token_pool(outputs.last_hidden_state, batch_dict['attention_mask'])
     embeddings = F.normalize(embeddings, p=2, dim=1)