Method size

tensorrt_llm/functional.py:480–490

Returns the shape of the tensor if the dim parameter is None. Otherwise, returns a size of the dimension indicated by dim. The behavior is undefined if dim is negative or exceeds the rank of the tensor.

(self, dim=None)

        return cast(self, dtype)

    def size(self, dim=None):
        '''
        Returns the shape of the tensor if the dim parameter is None.
        Otherwise, returns a size of the dimension indicated by dim. The
        behavior is undefined if dim is negative or exceeds the rank of the
        tensor.
        '''
        if dim is None:
            return self.trt_tensor.shape

        return self.trt_tensor.shape[dim]

    def rank(self):
        '''