(model: nn.Module)
| 46 | |
| 47 | |
| 48 | def get_weight_parallel_dim(model: nn.Module): |
| 49 | weight_parallel_dim = {} |
| 50 | for name, module in model.named_modules(): |
| 51 | for class_, dict_ in _MODEL_PARALLEL_MODULES: |
| 52 | if isinstance(module, class_): |
| 53 | for leaf_name, dim in dict_.items(): |
| 54 | full_name = name + "." + leaf_name if name else leaf_name |
| 55 | if dim >= 0: |
| 56 | weight_parallel_dim[full_name] = dim |
| 57 | break |
| 58 | return weight_parallel_dim |
| 59 | |
| 60 | |
| 61 | def _tensor_list_max_diff(tensors: List[torch.Tensor]) -> float: |
no outgoing calls
no test coverage detected