如何为每一轮将pytorch模型输出存储到numpy

Question

我想为每一轮将pytorch模型的预测输出存储到numpy。因为我使用ResNet18作为主干模型，4000个图像样本作为输入。它说RuntimeError: CUDA out of memory. Tried to allocate 200.00 MiB (GPU 0; 10.00 GiB total capacity; 7.50 GiB already allocated; 0 bytes free; 7.80 GiB reserved in total by PyTorch)，我发现这可能是由于批量太大造成的。如何将预测输出存储到批量大小有限的numpy

    model.eval()
    for i in range(len(dataset)):
        x = torch.Tensor(dataset[i]['image']).to(device)
        predict = model(x)
        predict_prob = torch.softmax(predict, dim=1).cpu()
        dataset[i]['label'] = predict_prob.detach().numpy()