From 92cda1b7da41551b2c177456079b5d77323f082d Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Mon, 5 Aug 2024 06:41:12 +0000
Subject: [PATCH] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 scripts/infer.py | 33 +++++++++++++++++++--------------
 1 file changed, 19 insertions(+), 14 deletions(-)

diff --git a/scripts/infer.py b/scripts/infer.py
index e542d46..6043b98 100644
--- a/scripts/infer.py
+++ b/scripts/infer.py
@@ -35,9 +35,10 @@
 
 try:
     from monai.utils import TRTWrapper
-    TRT_AVAILABLE=True
-except Exception as e:
-    TRT_AVAILABLE=False
+
+    TRT_AVAILABLE = True
+except Exception:
+    TRT_AVAILABLE = False
 
 rearrange, _ = optional_import("einops", name="rearrange")
 sys.path.insert(0, os.path.abspath(os.path.dirname(__file__)))
@@ -137,19 +138,23 @@ def __init__(self, config_file="./configs/infer.yaml", **override):
         self.prev_mask = None
         self.batch_data = None
         if self.trt and TRT_AVAILABLE:
-            ts=os.path.getmtime(config_file)
-            self.model.image_encoder.encoder = TRTWrapper("Encoder",
-                                                          self.model.image_encoder.encoder,
-                                                          input_names=["x"],
-                                                          output_names=["x_out"],
-                                                          timestamp=ts)
+            ts = os.path.getmtime(config_file)
+            self.model.image_encoder.encoder = TRTWrapper(
+                "Encoder",
+                self.model.image_encoder.encoder,
+                input_names=["x"],
+                output_names=["x_out"],
+                timestamp=ts,
+            )
             self.model.image_encoder.encoder.load_engine()
 
-            self.model.class_head = TRTWrapper("ClassHead",
-                                               self.model.class_head,  
-                                               input_names=["src", "class_vector"],
-                                               output_names=["masks", "class_embedding"],
-                                               timestamp=ts)
+            self.model.class_head = TRTWrapper(
+                "ClassHead",
+                self.model.class_head,
+                input_names=["src", "class_vector"],
+                output_names=["masks", "class_embedding"],
+                timestamp=ts,
+            )
             self.model.class_head.load_engine()
         return