Disables grads in extract tool

moabitcoin · Oct 15, 2019 · 686f105 · 686f105
1 parent 14e0a01
commit 686f105
Showing 1 changed file with 7 additions and 6 deletions.
diff --git a/ig65m/cli/extract.py b/ig65m/cli/extract.py
@@ -67,14 +67,15 @@ def main(args):
 
     features = []
 
-    for inputs in tqdm(loader, total=len(dataset) // args.batch_size):
-        inputs = inputs.to(device)
+    with torch.no_grad():
+        for inputs in tqdm(loader, total=len(dataset) // args.batch_size):
+            inputs = inputs.to(device)
 
-        outputs = model(inputs)
-        outputs = outputs.data.cpu().numpy()
+            outputs = model(inputs)
+            outputs = outputs.data.cpu().numpy()
 
-        for output in outputs:
-            features.append(output)
+            for output in outputs:
+                features.append(output)
 
     np.save(args.features, np.array(features), allow_pickle=False)
     print("🍪 Done", file=sys.stderr)