File size: 400 Bytes
cc3cebf |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 |
"""module for gpu capabilities"""
from typing import Optional
from pydantic import BaseModel, Field
class GPUCapabilities(BaseModel):
"""model to manage the gpu capabilities statically"""
bf16: bool = Field(default=False)
fp8: bool = Field(default=False)
n_gpu: int = Field(default=1)
n_node: int = Field(default=1)
compute_capability: Optional[str] = Field(default=None)
|