r""" Move internal buffers of the ExponentialMovingAverage to `device`. Args: device: like `device` argument to `torch.Tensor.to`
(self, device=None, dtype=None, non_blocking=False)
| 787 | self.shadow_params = [p.pin_memory() for p in self.shadow_params] |
| 788 | |
| 789 | def to(self, device=None, dtype=None, non_blocking=False) -> None: |
| 790 | r""" |
| 791 | Move internal buffers of the ExponentialMovingAverage to `device`. |
| 792 | |
| 793 | Args: |
| 794 | device: like `device` argument to `torch.Tensor.to` |
| 795 | """ |
| 796 | # .to() on the tensors handles None correctly |
| 797 | self.shadow_params = [ |
| 798 | p.to(device=device, dtype=dtype, non_blocking=non_blocking) |
| 799 | if p.is_floating_point() |
| 800 | else p.to(device=device, non_blocking=non_blocking) |
| 801 | for p in self.shadow_params |
| 802 | ] |
| 803 | |
| 804 | def state_dict(self) -> dict: |
| 805 | r""" |
no outgoing calls
no test coverage detected