Load dolly-v2 model with remote code trusted and full text returned (so it's usable with langchain)

handler.py

@@ -1,3 +1,4 @@
+import torch
 from typing import Dict, Any, List
 
 from transformers import pipeline
@@ -8,7 +9,7 @@ class EndpointHandler:
             self,
             path: str,
     ) -> None:
-        self.pipeline = pipeline(model=path, trust_remote_code=True, return_full_text=True)
+        self.pipeline = pipeline(model=path, torch_dtype=torch.bfloat16, trust_remote_code=True, return_full_text=True)
 
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         """