| |
| def _type(self, new_type=None, async=False): |
| if new_type is None: |
| return self.__module__ + '.' + self.__class__.__name__ |
| |
| if isinstance(new_type, str): |
| new_type = _import_dotted_name(new_type) |
| if new_type == type(self): |
| return self |
| return new_type(self.size()).copy_(self, async) |
| |
| def _cuda(self, idx=None, async=False): |
| import torch.cuda |
| # This already is a CUDA tensor. |
| # Let's check if it needs to be transfered to another GPU. |
| if hasattr(self, 'get_device'): |
| target_device = idx if idx else torch.cuda.current_device() |
| if self.get_device() != target_device: |
| with torch.cuda.device(target_device): |
| return type(self)(self.size()).copy_(self, async) |
| else: |
| return self |
| else: |
| ctx = torch.cuda.device(idx) if idx else torch.cuda._dummy_ctx() |
| with ctx: |
| return self.type(getattr(torch.cuda, self.__class__.__name__), async) |
| |
| def _range(*args, **kwargs): |
| return __builtins__['range'](*args, **kwargs) |
| |
| |
| def _import_dotted_name(name): |
| components = name.split('.') |
| obj = __import__(components[0]) |
| for component in components[1:]: |
| obj = getattr(obj, component) |
| return obj |