disable pin memory + bug fix

mohamad-hasan-sohan-ajini · Jul 9, 2020 · 0f1cf5d · 0f1cf5d
1 parent ace014f
commit 0f1cf5d
Showing 1 changed file with 3 additions and 4 deletions.
diff --git a/trainer.py b/trainer.py
@@ -50,8 +50,7 @@
     ds,
     batch_size=args.batch_size,
     shuffle=True,
-    num_workers=args.num_workers,
-    pin_memory=True
+    num_workers=args.num_workers
 )
 feature_extractor = transform(**kwargs).to(device)
 
@@ -92,7 +91,7 @@
         with torch.no_grad():
             x = x.to(device)
             x = feature_extractor(x) + 1
-            x = x.log()
+            x = x.log().detach()
         target = target.to(device)
 
         # forward pass
@@ -134,4 +133,4 @@
         if args.save_checkpoint:
             save_checkpoint(args.save_path, model, criterion, optimizer, epoch)
 
-save_checkpoint(model, criterion, optimizer, epoch)
+save_checkpoint(args.save_path, model, criterion, optimizer, epoch)