gaspersic/triton
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[FRONTEND] updated TensorWrapper (#299)

Browse Source
This commit is contained in:
Philippe Tillet
2021-09-22 13:53:27 -07:00
committed by GitHub
parent 2849e7a773
commit 5211f23a63
1 changed files with 7 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
python/triton/code_gen.py
View File

@@ -936,16 +936,16 @@ def next_power_of_2(n):
###### ######
class TensorWrapper: class TensorWrapper:
def __init__(self, data_ptr, dtype, device): def __init__(self, base, dtype):
self._data_ptr = data_ptr
self.dtype = dtype self.dtype = dtype
self.device = device self.base = base
self.is_cuda = base.is_cuda
self.device = base.device
def data_ptr(self): def data_ptr(self):
return self._data_ptr return self.base.data_ptr()
def reinterpret(tensor, dtype): def reinterpret(tensor, dtype):
return TensorWrapper(tensor.data_ptr(), dtype, tensor.device) return TensorWrapper(tensor, dtype)