fix: add missing imports

server/lorax_server/utils/adapter.py

@@ -4,6 +4,7 @@
 from typing import TYPE_CHECKING, Set, Tuple, Union
 
 import torch
+from safetensors import safe_open
 from safetensors.torch import load_file
 from transformers import AutoConfig, AutoTokenizer, PreTrainedTokenizer
 
@@ -166,4 +167,4 @@ def load_module_weight(name: str, module: Union[torch.Tensor, str], device, dtyp
 
  # module would be just the filename if lazy loading happened before
  with safe_open(module, framework="pt", device=device) as f:
- return f.get_tensor(name).to(dtype)
+ return f.get_tensor(name).to(dtype)