"""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)