| 77 | |
| 78 | |
| 79 | def get_parameter_device(parameter: torch.nn.Module) -> torch.device: |
| 80 | try: |
| 81 | parameters_and_buffers = itertools.chain(parameter.parameters(), parameter.buffers()) |
| 82 | return next(parameters_and_buffers).device |
| 83 | except StopIteration: |
| 84 | # For torch.nn.DataParallel compatibility in PyTorch 1.5 |
| 85 | |
| 86 | def find_tensor_attributes(module: torch.nn.Module) -> List[Tuple[str, Tensor]]: |
| 87 | tuples = [(k, v) for k, v in module.__dict__.items() if torch.is_tensor(v)] |
| 88 | return tuples |
| 89 | |
| 90 | gen = parameter._named_members(get_members_fn=find_tensor_attributes) |
| 91 | first_tuple = next(gen) |
| 92 | return first_tuple[1].device |
| 93 | |
| 94 | |
| 95 | def get_parameter_dtype(parameter: torch.nn.Module) -> torch.dtype: |