import torch
..
from dropbp.layer import DropBP

...
class Block(nn.Modoule): # transformer block
    def __init__(self, ..):
        self.norm_1 = ...
        self.attn = ...
        self.norm_2 = ...
        self.mlp = ...
        
        # Define DropBP layers
        # The FLOPs below is about general transformer block per batch*seq
        # with intermediate_size = 4*hidden_size
        attn_flops = 8*config.hidden_size**2 + 4*config.hidden_size*self.sequence_length 
        mlp_flops = 16*config.hidden_size**2
        self.dropbp_attn = DropBP(flops=attn_flops)
        self.dropbp_mlp = DropBP(flops=mlp_flops)
        ...
    def forward(self, x, ..):
        h = self.attn(self.norm_1(x), ...)
        x = self.dropbp_attn(h)+x   # instead of 'x = h+x'  
        h = self.mlp(self.norm_2(x))
        x = self.dropbp_mlp(h)+x    # instead of 'x = h+x'    
        return x

import torch
from dropbp.handler import DropBPHandler

model = ... # user define model
optimizer = ... # user define optimizer

dropbp_handler = DropBPHandler(model) # define drop handler
dropbp_handler.set_initial_drop_rate(drop_rate) # set a drop rate

# training loop
for iter in ...
    def backprop: # define backprop
        output = model(data)
        loss = loss_func(output, target)
        optimizer.zero_grad() # this line must be present
        loss.backward()

    if iter == int(max_iter * 0.1) # adjust drop rates at 10% of training process 
        dropbp_handler.sensitivity_based_drop_bp(backprop, drop_rate) # it automatically adjusts drop rates
    
    out = model(data)
    loss = loss_func(output,target)
    non_grad = dropbp_handler.detact_non_grad() # detect when all layers are dropped
    if not(non_grad): # exclude the above situation for avoiding error
        loss.backward()
    optimizer.step()
    ...