bug描述 Describe the Bug

• paddle版本develop。

在jupyter notebook中初始化一下这个tensor！矩阵太大了需要分配一段时间。

import paddle, time
tensor = paddle.randn((1, 4096, 1024))

下面的代码在jupyter notebook中循环点击运行N次，发现结果正常。

def logsumexp(x, axis=0, keepdim=False):
    return paddle.log(x.exp().sum(axis=axis, keepdim=keepdim))

start = time.time()
for i in range(100):
    logsumexp(tensor, axis=1, keepdim=True)
end = time.time()
print("custom logsumexp", end-start)

下面的代码在jupyter notebook中循环点击运行N次，发现程序卡住了？感觉好像卡在显存分配上？？？？前几步训练还正常，但是你循环运行N次后代码就挂了。

start = time.time()
for i in range(100):
    paddle.logsumexp(tensor, axis=1, keepdim=True)
end = time.time()
print("paddle logsumexp", end-start)

其他补充信息 Additional Supplementary Information

No response