import torch
import torch.nn as nn

class BitFit:
    def __init__(self, model):
        self.model = model
        self.trainable_count = 0

    def wrap(self):
        for p in self.model.parameters():
            p.requires_grad = False
        for name, param in self.model.named_parameters():
            if 'bias' in name:
                param.data = param.data.to(device=param.device, dtype=param.dtype)
                param.requires_grad = True
                self.trainable_count += 1

    def parameters(self):
        for n, p in self.model.named_parameters():
            if p.requires_grad:
                yield p

def inject_bitfit(model):
    peft = BitFit(model)
    peft.wrap()
    return model, peft