mirror of
https://github.com/wassname/vllm.git
synced 2026-07-20 12:40:26 +08:00
18 lines
378 B
Python
18 lines
378 B
Python
from typing import Tuple
|
|
|
|
import torch
|
|
|
|
from .interface import Platform, PlatformEnum
|
|
|
|
|
|
class TpuPlatform(Platform):
|
|
_enum = PlatformEnum.TPU
|
|
|
|
@staticmethod
|
|
def get_device_capability(device_id: int = 0) -> Tuple[int, int]:
|
|
raise RuntimeError("TPU does not have device capability.")
|
|
|
|
@staticmethod
|
|
def inference_mode():
|
|
return torch.no_grad()
|