Get torch.device from module, assuming that the whole module has one device.
(self)
| 125 | |
| 126 | @property |
| 127 | def device(self) -> device: |
| 128 | """ |
| 129 | Get torch.device from module, assuming that the whole module has one device. |
| 130 | """ |
| 131 | try: |
| 132 | return next(self.parameters()).device |
| 133 | except StopIteration: |
| 134 | # For nn.DataParallel compatibility in PyTorch 1.5 |
| 135 | |
| 136 | def find_tensor_attributes(module: nn.Module) -> List[Tuple[str, Tensor]]: |
| 137 | tuples = [(k, v) for k, v in module.__dict__.items() if torch.is_tensor(v)] |
| 138 | return tuples |
| 139 | |
| 140 | gen = self._named_members(get_members_fn=find_tensor_attributes) |
| 141 | first_tuple = next(gen) |
| 142 | return first_tuple[1].device |
| 143 | |
| 144 | @property |
| 145 | def dtype(self) -> dtype: |
nothing calls this directly
no outgoing calls
no test coverage detected