def __exit__(self, *args): if self.prev_idx != self.idx: torch._C._cuda_setDevice(self.prev_idx) return False