utils.py

import torch

def _get_compute_dtype():
    major, minor = torch.cuda.get_device_capability()
    if major >= 8:
        return torch.bfloat16
    else:
        return torch.float16


def print_trainable_parameters(model):
    """
    Prints the number of trainable parameters in the model.
    """
    trainable_params = 0
    all_param = 0
    for _, param in model.named_parameters():
        all_param += param.numel()
        if param.requires_grad:
            trainable_params += param.numel()
    print(
        f"trainable params: {trainable_params} || "
        f"all params: {all_param} || "
        f"trainable: {100 * trainable_params / all_param}"
    )