line 77: add torch.no_grad()

Add torch.no_grad() to reduce memory consumption during inference
wctu · Apr 4, 2019 · 66317a9 · 66317a9
1 parent 17584dc
commit 66317a9
Showing 1 changed file with 6 additions and 5 deletions.
diff --git a/test/test.py b/test/test.py
@@ -74,11 +74,12 @@ def main():
         inputs = inputs.cuda(gpu_id, non_blocking=True)
 
         # inference
-        out_x = model(inputs)
-        inputs_t = torch.transpose(inputs, 2, 3)
-        out_y_t = model(inputs_t)
-        out_y = torch.transpose(out_y_t, 2, 3)
-        outputs = torch.cat((out_x, out_y), 1)
+        with torch.no_grad():
+            out_x = model(inputs)
+            inputs_t = torch.transpose(inputs, 2, 3)
+            out_y_t = model(inputs_t)
+            out_y = torch.transpose(out_y_t, 2, 3)
+            outputs = torch.cat((out_x, out_y), 1)
 
         # compute superpixels
         affinity = outputs[0].data.cpu().numpy()