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

[PYTHON] Update TensorWrapper with device attribute (#102)

Browse Source 
bugfix `triton.reinterpret`
This commit is contained in:
Szymon Sidor
2021-05-07 23:51:02 -07:00
committed by Philippe Tillet
parent 840140bf26
commit ac57812bdc
1 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
python/triton/code_gen.py
View File

@@ -711,15 +711,15 @@ def cdiv(x, y):
###### ######
class TensorWrapper: class TensorWrapper:
def __init__(self, data_ptr, dtype): def __init__(self, data_ptr, dtype, device):
self._data_ptr = data_ptr self._data_ptr = data_ptr
self.dtype = dtype self.dtype = dtype
self.device = device
def data_ptr(self): def data_ptr(self):
return self._data_ptr return self._data_ptr
def reinterpret(tensor, dtype): def reinterpret(tensor, dtype):
return TensorWrapper(tensor.data_ptr(), dtype) return TensorWrapper(tensor.data_ptr(), dtype, tensor.device)