So my workaround was to use the per-layer learning rates and use one weight decay value for all the parameters.
optimizer_parameters = [
#         {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.001},
#         {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0},
        {'params': model.roberta.parameters(), 'lr': lr[0]},        
        {'params': model.last_linear.parameters(), 'lr': lr[1]}    
        
    ]

    optimizer = AdamW(optimizer_parameters, weight_decay=0.001)

Would that affect performance since now I’m applying weight decay to the bias and LayerNorm params, or is it ok?
And thanks btw!