code quality edits

neuralmagic · Jan 27, 2023 · 2ae697b · 2ae697b
1 parent d0a6523
commit 2ae697b
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 7 deletions.
diff --git a/src/sparseml/pytorch/torchvision/train.py b/src/sparseml/pytorch/torchvision/train.py
@@ -139,13 +139,17 @@ def train_one_epoch(
                 # Reset ema buffer to keep copying weights during warmup period
                 model_ema.n_averaged.fill_(0)
 
-        acc1, num_correct_1, acc5, num_correct_5 = utils.accuracy(output, target, topk=(1, 5))
+        acc1, num_correct_1, acc5, num_correct_5 = utils.accuracy(
+            output, target, topk=(1, 5)
+        )
         batch_size = image.shape[0]
         metric_logger.update(loss=loss.item(), lr=optimizer.param_groups[0]["lr"])
         metric_logger.meters["acc1"].update(
-                acc1.item(), n=batch_size, total=num_correct_1)
+            acc1.item(), n=batch_size, total=num_correct_1
+        )
         metric_logger.meters["acc5"].update(
-            acc5.item(), n=batch_size, total=num_correct_5)
+            acc5.item(), n=batch_size, total=num_correct_5
+        )
         metric_logger.meters["imgs_per_sec"].update(
             batch_size / (time.time() - start_time)
         )
@@ -184,15 +188,18 @@ def evaluate(
             loss = criterion(output, target)
 
             acc1, num_correct_1, acc5, num_correct_5 = utils.accuracy(
-                output, target, topk=(1, 5))
+                output, target, topk=(1, 5)
+            )
             # FIXME need to take into account that the datasets
             # could have been padded in distributed setup
             batch_size = image.shape[0]
             metric_logger.update(loss=loss.item())
             metric_logger.meters["acc1"].update(
-                acc1.item(), n=batch_size, total=num_correct_1)
+                acc1.item(), n=batch_size, total=num_correct_1
+            )
             metric_logger.meters["acc5"].update(
-                acc5.item(), n=batch_size, total=num_correct_5)
+                acc5.item(), n=batch_size, total=num_correct_5
+            )
             num_processed_samples += batch_size
     # gather the stats from all processes
 

diff --git a/src/sparseml/pytorch/torchvision/utils.py b/src/sparseml/pytorch/torchvision/utils.py
@@ -213,7 +213,7 @@ def accuracy(output, target, topk=(1,)):
         for k in topk:
             correct_k = correct[:k].flatten().sum(dtype=torch.float32)
             res.append(correct_k * (100.0 / batch_size))
-            res.append(correct_k *  100.0)
+            res.append(correct_k * 100.0)
         return res