| import torch | |
| def rewrite_model(model, plan): | |
| for name,module in model.named_modules(): | |
| if name in plan: | |
| precision=plan[name] | |
| for pname,p in module.named_parameters(recurse=False): | |
| if precision=="fp16": | |
| p.data = p.data.half() | |
| else: | |
| p.data = p.data.float() | |
| return model | |