From 03be9ff21368894aa82c89c4813862763afa1327 Mon Sep 17 00:00:00 2001
From: wengkaiheng <wengkaiheng@meituan.com>
Date: Fri, 8 Jul 2022 01:02:40 +0800
Subject: [PATCH] fix: update minor

---
 yolov6/core/engine.py           |  2 +-
 yolov6/core/evaler.py           |  2 +-
 yolov6/layers/dbb_transforms.py |  2 +-
 yolov6/models/end2end.py        | 24 +++++++++++++-----------
 4 files changed, 16 insertions(+), 14 deletions(-)

diff --git a/yolov6/core/engine.py b/yolov6/core/engine.py
index 0b9f5dd9..879399bc 100644
--- a/yolov6/core/engine.py
+++ b/yolov6/core/engine.py
@@ -22,6 +22,7 @@
 from yolov6.utils.checkpoint import load_state_dict, save_checkpoint, strip_optimizer
 from yolov6.solver.build import build_optimizer, build_lr_scheduler
 
+
 class Trainer:
     def __init__(self, args, cfg, device):
         self.args = args
@@ -65,7 +66,6 @@ def __init__(self, args, cfg, device):
         self.batch_size = args.batch_size
         self.img_size = args.img_size
 
-
     # Training Process
 
     def train(self):
diff --git a/yolov6/core/evaler.py b/yolov6/core/evaler.py
index 569e4e3b..bd4453da 100644
--- a/yolov6/core/evaler.py
+++ b/yolov6/core/evaler.py
@@ -207,7 +207,7 @@ def convert_to_coco_format(self, outputs, imgs, paths, shapes, ids):
 
     @staticmethod
     def check_task(task):
-        if task not in ['train','val','speed']:
+        if task not in ['train', 'val', 'speed']:
             raise Exception("task argument error: only support 'train' / 'val' / 'speed' task.")
 
     @staticmethod
diff --git a/yolov6/layers/dbb_transforms.py b/yolov6/layers/dbb_transforms.py
index e60cbd4d..cd93d0e2 100644
--- a/yolov6/layers/dbb_transforms.py
+++ b/yolov6/layers/dbb_transforms.py
@@ -27,7 +27,7 @@ def transIII_1x1_kxk(k1, b1, k2, b2, groups):
             k1_T_slice = k1_T[:, g*k1_group_width:(g+1)*k1_group_width, :, :]
             k2_slice = k2[g*k2_group_width:(g+1)*k2_group_width, :, :, :]
             k_slices.append(F.conv2d(k2_slice, k1_T_slice))
-            b_slices.append((k2_slice * b1[g*k1_group_width:(g+1)*k1_group_width].reshape(1, -1, 1, 1)).sum((1, 2, 3)))
+            b_slices.append((k2_slice * b1[g * k1_group_width:(g+1) * k1_group_width].reshape(1, -1, 1, 1)).sum((1, 2, 3)))
         k, b_hat = transIV_depthconcat(k_slices, b_slices)
     return k, b_hat + b2
 
diff --git a/yolov6/models/end2end.py b/yolov6/models/end2end.py
index ee048e24..e3902b57 100644
--- a/yolov6/models/end2end.py
+++ b/yolov6/models/end2end.py
@@ -2,6 +2,7 @@
 import torch.nn as nn
 import random
 
+
 class ORT_NMS(torch.autograd.Function):
 
     @staticmethod
@@ -25,6 +26,7 @@ def forward(ctx,
     def symbolic(g, boxes, scores, max_output_boxes_per_class, iou_threshold, score_threshold):
         return g.op("NonMaxSuppression", boxes, scores, max_output_boxes_per_class, iou_threshold, score_threshold)
 
+
 class TRT_NMS(torch.autograd.Function):
     @staticmethod
     def forward(
@@ -59,18 +61,18 @@ def symbolic(g,
                  score_activation=0,
                  score_threshold=0.25):
         out = g.op("TRT::EfficientNMS_TRT",
-                     boxes,
-                     scores,
-                     background_class_i=background_class,
-                     box_coding_i=box_coding,
-                     iou_threshold_f=iou_threshold,
-                     max_output_boxes_i=max_output_boxes,
-                     plugin_version_s=plugin_version,
-                     score_activation_i=score_activation,
-                     score_threshold_f=score_threshold,
-                     outputs=4)
+                   boxes,
+                   scores,
+                   background_class_i=background_class,
+                   box_coding_i=box_coding,
+                   iou_threshold_f=iou_threshold,
+                   max_output_boxes_i=max_output_boxes,
+                   plugin_version_s=plugin_version,
+                   score_activation_i=score_activation,
+                   score_threshold_f=score_threshold,
+                   outputs=4)
         nums, boxes, scores, classes = out
-        return nums,boxes,scores,classes
+        return nums, boxes, scores, classes