Set torch default dtype in a context manager (#971)

Signed-off-by: Antoni Baum <antoni.baum@protonmail.com>

vllm/model_executor/model_loader.py

 """Utilities for selecting and loading models."""
+import contextlib
 from typing import Type
 
 import torch
@@ -30,6 +31,15 @@ _MODEL_REGISTRY = {
 }
 
 
+@contextlib.contextmanager
+def _set_default_torch_dtype(dtype: torch.dtype):
+    """Sets the default torch dtype to the given dtype."""
+    old_dtype = torch.get_default_dtype()
+    torch.set_default_dtype(dtype)
+    yield
+    torch.set_default_dtype(old_dtype)
+
+
 def _get_model_architecture(config: PretrainedConfig) -> Type[nn.Module]:
     architectures = getattr(config, "architectures", [])
     for arch in architectures:
@@ -42,8 +52,7 @@ def _get_model_architecture(config: PretrainedConfig) -> Type[nn.Module]:
 
 def get_model(model_config: ModelConfig) -> nn.Module:
     model_class = _get_model_architecture(model_config.hf_config)
-    torch.set_default_dtype(model_config.dtype)
-
+    with _set_default_torch_dtype(model_config.dtype):
         # Create a model instance.
         # The weights will be initialized as empty tensors.
         model = model_class(model_config.hf_config)