Merge pull request #439 from NVIDIA/clean_up_ptq_recipe

refactor: Fix python linting issues
pytorch · Apr 21, 2021 · c3b8583 · c3b8583
2 parents 2af2c11 + a39dea7
commit c3b8583
Show file tree

Hide file tree

Showing 5 changed files with 5 additions and 3 deletions.
diff --git a/.github/scripts/run_cpp_linter.py b/.github/scripts/run_cpp_linter.py
@@ -23,6 +23,7 @@
 if output.returncode != 0:
     comment = '''There are some changes that do not conform to C++ style guidelines:\n ```diff\n{}```'''.format(output.stdout.decode("utf-8"))
     approval = 'REQUEST_CHANGES'
+    exit(1)
 
 pr.create_review(commit, comment, approval)
 

diff --git a/.github/scripts/run_py_linter.py b/.github/scripts/run_py_linter.py
@@ -23,5 +23,6 @@
 if output.returncode != 0:
     comment = '''There are some changes that do not conform to Python style guidelines:\n ```diff\n{}```'''.format(output.stdout.decode("utf-8"))
     approval = 'REQUEST_CHANGES'
+    exit(1)
 
 pr.create_review(commit, comment, approval)
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"]