From 7385f3d20eb90f03a0b0b30379fc56aa638aaa39 Mon Sep 17 00:00:00 2001 From: Naravich Chutisilp Date: Sun, 15 Nov 2020 20:13:19 +0700 Subject: [PATCH] Fix "RuntimeError: CUDA Error: out of memory" --- model.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/model.py b/model.py index 8bee563..9793ae3 100644 --- a/model.py +++ b/model.py @@ -72,9 +72,10 @@ def main(): gt = torch.cat((gt, target), 0) bs, n_crops, c, h, w = inp.size() input_var = torch.autograd.Variable(inp.view(-1, c, h, w).cuda(), volatile=True) - output = model(input_var) - output_mean = output.view(bs, n_crops, -1).mean(1) - pred = torch.cat((pred, output_mean.data), 0) + with torch.no_grad(): # Fix CUDA out of memory + output = model(input_var) + output_mean = output.view(bs, n_crops, -1).mean(1) + pred = torch.cat((pred, output_mean.data), 0) AUROCs = compute_AUCs(gt, pred) AUROC_avg = np.array(AUROCs).mean() @@ -126,4 +127,4 @@ def forward(self, x): if __name__ == '__main__': - main() \ No newline at end of file + main()