[MISC] fix pin_memory=torch.cuda.is_available(), use is_pin_memory_available (#37415)
Signed-off-by: Kunshang Ji <kunshang.ji@intel.com>
This commit is contained in:
@@ -11,6 +11,7 @@ from transformers import PreTrainedTokenizerBase
|
||||
|
||||
from vllm.sampling_params import SamplingParams
|
||||
from vllm.utils.import_utils import LazyLoader
|
||||
from vllm.utils.platform_utils import is_pin_memory_available
|
||||
from vllm.v1.structured_output.backend_types import (
|
||||
StructuredOutputBackend,
|
||||
StructuredOutputGrammar,
|
||||
@@ -138,7 +139,7 @@ class LMFormatEnforcerBackend(StructuredOutputBackend):
|
||||
(max_num_seqs, (self.vocab_size + 31) // 32),
|
||||
-1,
|
||||
dtype=torch.int32,
|
||||
pin_memory=torch.cuda.is_available(),
|
||||
pin_memory=is_pin_memory_available(),
|
||||
)
|
||||
|
||||
def destroy(self):
|
||||
|
||||
@@ -15,6 +15,7 @@ from regex import escape as regex_escape
|
||||
|
||||
from vllm.sampling_params import SamplingParams
|
||||
from vllm.utils.import_utils import LazyLoader
|
||||
from vllm.utils.platform_utils import is_pin_memory_available
|
||||
from vllm.v1.structured_output.backend_types import (
|
||||
StructuredOutputBackend,
|
||||
StructuredOutputGrammar,
|
||||
@@ -96,7 +97,7 @@ class OutlinesBackend(StructuredOutputBackend):
|
||||
(max_num_seqs, (self.vocab_size + 31) // 32),
|
||||
-1,
|
||||
dtype=torch.int32,
|
||||
pin_memory=torch.cuda.is_available(),
|
||||
pin_memory=is_pin_memory_available(),
|
||||
)
|
||||
|
||||
def destroy(self):
|
||||
|
||||
Reference in New Issue
Block a user