时间戳

首页 > 解决方案 > PyTorch GPU 内存管理

pytorch - PyTorch GPU 内存管理

问题描述

在我的代码中，我想替换张量中的值，因为某些索引的值为零，例如

target_mac_out[avail_actions[:, 1:] == 0] = -9999999

但是，它返回 OOM

RuntimeError: CUDA out of memory. Tried to allocate 166.00 MiB (GPU 0; 10.76 GiB total capacity; 9.45 GiB already allocated; 4.75 MiB free; 9.71 GiB reserved in total by PyTorch)

我认为没有内存分配，因为它只是访问张量target_mac_out并检查值并替换某些索引的新值。

我理解对了吗？

标签： pytorch

解决方案

创建一个新的avail_actions[:, 1:] == 0张量，并且可能整行本身在完成操作后删除旧张量之前创建另一个张量。

如果速度不是问题，那么您可以使用for循环。喜欢

for i in range(target_mac_out.size(0)):
    for j in range(target_mac_out.size(1)-1):
        if target_mac_out[i, j+1] == 0:
            target_mac_out[i, j+1] = -9999999

推荐阅读