    def size(self) -> int:
        """
        Override size() to avoid D2H copy.
        """
        return self.tensor.element_size() * self.tensor.nelement()
Override size() to avoid D2H copy.
python/aitemplate/compiler/base.py