configs.py

#GPT2 Params 
from dataclasses import dataclass

@dataclass
class ModelConfig:
    block_size: int = 2048
    n_embd : int = 768
    n_head : int = 12
    n_layer : int = 12
    multiple_of : int = 256
    vocab_size : int = 50304
    device : str = "cuda"

    #Needed for Kv Cahce
    max_batch_size : int = 8