Move task to device/dtype, keeping self.device and self.dtype in sync.
(self, *args, **kwargs)
| 54 | self.verbose = verbose |
| 55 | |
| 56 | def to(self, *args, **kwargs): |
| 57 | """Move task to device/dtype, keeping self.device and self.dtype in sync.""" |
| 58 | dummy = torch.empty(0).to(*args, **kwargs) |
| 59 | self.device = dummy.device |
| 60 | self.dtype = dummy.dtype |
| 61 | return super().to(*args, **kwargs) |
| 62 | |
| 63 | def prepare_distributed( |
| 64 | self, |
no outgoing calls