def cuda( self, device: Optional[Union[torch_device, str, int]] = None) -> 'BaseTile': """Return a copy of this tile in CUDA memory. Args: device: CUDA device Returns: Self with the underlying C++ tile moved to CUDA memory. Raises: CudaError: if the library has not been compiled with CUDA. """ if not cuda.is_compiled(): raise CudaError('aihwkit has not been compiled with CUDA support') device = torch_device('cuda', cuda_device(device).idx) if self.is_cuda and device != self.device: raise CudaError( 'Cannot switch CUDA devices of existing Cuda tiles') if isinstance(self.tile, tiles.AnalogTile): with cuda_device(device): self.tile = tiles.CudaAnalogTile(self.tile) self.is_cuda = True self.device = device self.analog_ctx.cuda(device) return self
def cuda( self, device: Optional[Union[torch_device, str, int]] = None) -> BaseTile: """Return a copy of this tile in CUDA memory. Args: device: CUDA device """ if not cuda.is_compiled(): raise RuntimeError( 'aihwkit has not been compiled with CUDA support') with cuda_device(device): tile = CudaInferenceTile(self) # need also to copy construct! tile.alpha = self.alpha.cuda(device) if self.reference_combined_weights is not None: tile.reference_combined_weights = self.reference_combined_weights.to( device) if self.programmed_weights is not None: tile.programmed_weights = self.programmed_weights.to(device) if self.nu_drift_list is not None: tile.nu_drift_list = [nu.to(device) for nu in self.nu_drift_list] return tile
def cuda( self, device: Optional[Union[torch_device, str, int]] = None ) -> 'BaseTile': """Return a copy of this tile in CUDA memory.""" if not cuda.is_compiled(): raise RuntimeError('aihwkit has not been compiled with CUDA support') with cuda_device(device): tile = CudaFloatingPointTile(self) return tile
def cuda( self, device: Optional[Union[torch_device, str, int]] = None ) -> 'CudaIndexedAnalogTile': """Return a copy of this tile in CUDA memory. Args: device: CUDA device """ if not cuda.is_compiled(): raise RuntimeError( 'aihwkit has not been compiled with CUDA support') with cuda_device(device): tile = CudaIndexedAnalogTile(self) return tile
def cuda( self, device: Optional[Union[torch_device, str, int]] = None ) -> 'CudaIndexedFloatingPointTile': """Return a copy of this tile in CUDA memory. Args: device: CUDA device """ if not cuda.is_compiled(): raise RuntimeError( 'aihwkit has not been compiled with CUDA support') with cuda_device(device): tile = CudaIndexedFloatingPointTile(self.out_size, self.in_size, self.resistive_device, self.bias, self.in_trans, self.out_trans) return tile
def cuda( self, device: Optional[Union[torch_device, str, int]] = None ) -> 'BaseTile': """Return a copy of this tile in CUDA memory. Args: device: CUDA device Returns: A copy of this tile in CUDA memory. Raises: CudaError: if the library has not been compiled with CUDA. """ if not cuda.is_compiled(): raise CudaError('aihwkit has not been compiled with CUDA support') with cuda_device(device): tile = CudaFloatingPointTile(self) return tile