fix

for name, param in model.named_parameters():
    if name.startswith('head'): # fix MLP train backbone
        param.requires_grad = False

for name, param in model.named_parameters():
    if not name.startswith('head'): # fix backbone train MLP
        param.requires_grad = False