Skip to content

Commit

Permalink
Update
Browse files Browse the repository at this point in the history
  • Loading branch information
leimao committed May 12, 2024
1 parent 6dace6a commit a2b43b9
Showing 1 changed file with 20 additions and 4 deletions.
24 changes: 20 additions & 4 deletions python/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -74,10 +74,10 @@ def __init__(self,
(size, ))
self._device = cuda_call(cudart.cudaMalloc(nbytes))
self._nbytes = nbytes
self.name = name
self.shape = shape
self.format = format
self.dtype = dtype
self._name = name
self._shape = shape
self._format = format
self._dtype = dtype

@property
def host(self) -> np.ndarray:
Expand All @@ -99,6 +99,22 @@ def device(self) -> int:
def nbytes(self) -> int:
return self._nbytes

@property
def name(self) -> Optional[str]:
return self._name

@property
def shape(self) -> Optional[trt.Dims]:
return self._shape

@property
def format(self) -> Optional[trt.TensorFormat]:
return self._format

@property
def dtype(self) -> np.dtype:
return self._dtype

def __str__(self):
return f"Host:\n{self.host}\nDevice:\n{self.device}\nSize:\n{self.nbytes}\n"

Expand Down

0 comments on commit a2b43b9

Please sign in to comment.