refactor: Fix python linting issues

Signed-off-by: Naren Dasan <naren@narendasan.com> Signed-off-by: Naren Dasan <narens@nvidia.com>
pytorch · Apr 21, 2021 · 59868d2 · 59868d2
1 parent 2af2c11
commit 59868d2
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/cpp/ptq/training/vgg16/export_ckpt.py b/cpp/ptq/training/vgg16/export_ckpt.py
@@ -22,7 +22,7 @@ def test(model, dataloader, crit):
 
     with torch.no_grad():
         for data, labels in dataloader:
-            data, labels = data.cuda(), labels.cuda(async=True)
+            data, labels = data.cuda(), labels.cuda(non_blocking=True)
             out = model(data)
             loss += crit(out, labels)
             preds = torch.max(out, 1)[1]

diff --git a/cpp/ptq/training/vgg16/main.py b/cpp/ptq/training/vgg16/main.py
@@ -141,7 +141,7 @@ def train(model, dataloader, crit, opt, epoch):
     model.train()
     running_loss = 0.0
     for batch, (data, labels) in enumerate(dataloader):
-        data, labels = data.cuda(), labels.cuda(async=True)
+        data, labels = data.cuda(), labels.cuda(non_blocking=True)
         opt.zero_grad()
         out = model(data)
         loss = crit(out, labels)

diff --git a/py/trtorch/_compile_spec.py b/py/trtorch/_compile_spec.py
@@ -176,7 +176,7 @@ def _parse_compile_spec(compile_spec: Dict[str, Any]) -> trtorch._C.CompileSpec:
     if "max_batch_size" in compile_spec:
         assert type(compile_spec["max_batch_size"]) is int
         info.max_batch_size = compile_spec["max_batch_size"]
-    
+
     if "truncate_long_and_double" in compile_spec:
         assert type(compile_spec["truncate_long_and_double"]) is bool
         info.truncate_long_and_double = compile_spec["truncate_long_and_double"]