Optimised handler V2

handler.py

@@ -304,7 +304,7 @@ class EndpointHandler:
             with torch.no_grad():
                 for batch in text_loader:
                     batch_texts = batch[0].to(self.device, non_blocking=True)
-                    with torch.cuda.amp.autocast():
+                    with torch.amp.autocast():
                         embeddings = self.model(**batch_texts)
                     all_embeddings.append(embeddings)
             text_embeddings = torch.cat(all_embeddings, dim=0)
@@ -388,7 +388,7 @@ class EndpointHandler:
         if image_embeddings is not None and text_embeddings is not None:
             self.logger.info("Computing similarity scores.")
             try:
-                with torch.no_grad(), torch.cuda.amp.autocast():
+                with torch.no_grad(), torch.amp.autocast():
                     scores = self.processor.score_multi_vector(
                         text_embeddings, image_embeddings
                     )