Disable the redundant torch default initialization to accelerate model creation.
()
| 106 | |
| 107 | |
| 108 | def disable_torch_init(): |
| 109 | """ |
| 110 | Disable the redundant torch default initialization to accelerate model creation. |
| 111 | """ |
| 112 | import torch |
| 113 | |
| 114 | setattr(torch.nn.Linear, "reset_parameters", lambda self: None) |
| 115 | setattr(torch.nn.LayerNorm, "reset_parameters", lambda self: None) |
| 116 | |
| 117 | |
| 118 | def get_gpu_memory(max_gpus=None): |
nothing calls this directly
no outgoing calls
no test coverage detected