如何将张量分解成小块进行矩阵乘法以解决内存问题

Question

我有两个矩阵： 火炬尺寸（[1,8,50176,64]）火炬尺寸（[1,8,64,50176]） 我需要执行矩阵乘法来获得torch大小[1,8,50176,50176]，但我的cuda内存不足，需要75GB的GPU内存，我负担不起那么多GPU，我的GPU是8GB，我想知道如何将矩阵分解成小块并执行乘法，最后在一个torch张量中得到结果

我试过：

y = torch.matmul(a, b) 
y = torch.bmm(a.view(-1, 50176, 64), b.view(-1, 64, 50176)).view(1, 8, 50176, 50176)
y = torch.einsum('bdij, bdjk -> bdik', (a, b))

以上工作均未完成，均导致cuda内存不足错误

谢谢你的帮助

凌