"docs/backends/trtllm/gemma3_sliding_window_attention.md" did not exist on "b5fa0a34722ccf0b466c1306540a4669fda7890b"
tpu.py 307 Bytes
Newer Older
1
2
3
4
5
6
7
8
import torch

from .interface import Platform, PlatformEnum


class TpuPlatform(Platform):
    _enum = PlatformEnum.TPU

9
10
11
12
13
14
    @classmethod
    def get_device_name(cls, device_id: int = 0) -> str:
        raise NotImplementedError

    @classmethod
    def inference_mode(cls):
15
        return torch.no_grad()