diff --git a/transformers_server.py b/transformers_server.py index ebdb7af..905f9f5 100644 --- a/transformers_server.py +++ b/transformers_server.py @@ -79,7 +79,7 @@ def load_model(): processor_class = resolve_transformers_class(config["processer_class"]) model_class = resolve_transformers_class(config["model_class"]) torch_dtype = resolve_torch_dtype(config["torch_dtype"]) - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + device = torch.device("cuda") logger.info( f"model config: model_class={config['model_class']}, "