(module)
| 30 | |
| 31 | |
| 32 | def get_async_norm_states(module): |
| 33 | async_norm_states = OrderedDict() |
| 34 | for name, child in module.named_modules(): |
| 35 | if isinstance(child, ASYNC_NORM): |
| 36 | for k, v in child.state_dict().items(): |
| 37 | async_norm_states[".".join([name, k])] = v |
| 38 | return async_norm_states |
| 39 | |
| 40 | |
| 41 | def pyobj2tensor(pyobj, device="cuda"): |