(state, key)
| 1157 | """Move optimizer states to device. Note that this assumes the state structure of DeepSpeed Adam.""" |
| 1158 | |
| 1159 | def move_key(state, key): |
| 1160 | offload_buf_key = _make_offload_state_key(key) |
| 1161 | if offload_buf_key not in state: |
| 1162 | state[offload_buf_key] = torch.empty_like(state[key], device=device) |
| 1163 | if pin_memory: |
| 1164 | state[offload_buf_key] = get_accelerator().pin_memory(state[offload_buf_key]) |
| 1165 | state[offload_buf_key].copy_(state[key], non_blocking=non_blocking) |
| 1166 | state[key].data = state[offload_buf_key] |
| 1167 | |
| 1168 | for _, state in optimizer.state.items(): |
| 1169 | if "exp_avg" in state: |
no test coverage detected
searching dependent graphs…