Return a dictionary containing name to a shape and dtype.
(module: nn.Module)
| 56 | |
| 57 | |
| 58 | def get_weight_buffer_meta_from_module(module: nn.Module) -> Dict[str, Dict]: |
| 59 | """ |
| 60 | Return a dictionary containing name to a shape and dtype. |
| 61 | """ |
| 62 | weight_buffer_meta = {} |
| 63 | for name, param in sorted(module.named_parameters()): |
| 64 | weight_buffer_meta[name] = {'shape': param.shape, 'dtype': param.dtype} |
| 65 | return weight_buffer_meta |
| 66 | |
| 67 | |
| 68 | def build_memory_buffer(weight_buffer_meta: Dict[str, Dict]) -> Dict[torch.dtype, MemoryBuffer]: |
no test coverage detected