(tensor1: torch.Tensor, tensor2: torch.Tensor)
| 182 | |
| 183 | |
| 184 | def _get_diff_dim(tensor1: torch.Tensor, tensor2: torch.Tensor) -> int: |
| 185 | for i, (s1, s2) in enumerate(zip(tensor1.shape, tensor2.shape)): |
| 186 | if s1 != s2: |
| 187 | return i |
| 188 | return -1 |
| 189 | |
| 190 | |
| 191 | def _load_by_rank(checkpoint_dir: Union[str, bytes, os.PathLike], |
no outgoing calls
no test coverage detected
searching dependent graphs…