From 8e350c3c6e9fa9c8a051dd354f08d4a116bb39af Mon Sep 17 00:00:00 2001
From: clemente0731 <32806348+clemente0731@users.noreply.github.com>
Date: Thu, 4 Jul 2024 09:56:39 +0000
Subject: [PATCH] fix issue

---
 pprobe/bootstrap/hook_setup.py | 8 ++++----
 pprobe/toggle/cli.py           | 4 +---
 2 files changed, 5 insertions(+), 7 deletions(-)

diff --git a/pprobe/bootstrap/hook_setup.py b/pprobe/bootstrap/hook_setup.py
index 43ebb1b..5aa14f8 100644
--- a/pprobe/bootstrap/hook_setup.py
+++ b/pprobe/bootstrap/hook_setup.py
@@ -93,6 +93,9 @@ def check_and_run_hook(self, module_fullname):
                     pass
                 if self.torch_dump_op_enabled:
                     self.run_torch_func_hook()
+                if self.torch_dump_aten_enabled:
+                    # TODO
+                    pass
                 if self.torch_dump_dist_enabled:
                     # TODO
                     pass
@@ -101,15 +104,12 @@ def check_and_run_hook(self, module_fullname):
                 if self.torch_dump_optim_enabled:
                     # TODO
                     pass
-                if self.torch_dump_memory_enabled:
+                if self.torch_catch_memory_enabled:
                     # TODO
                     pass
                 if self.torch_test_dump_op_enabled:
                     # TODO
                     pass
-                if self.torch_test_dump_aten_enabled:
-                    # TODO
-                    pass
                 if self.torch_test_dump_dist_enabled:
                     # TODO
                     pass
diff --git a/pprobe/toggle/cli.py b/pprobe/toggle/cli.py
index 876467f..2607a2a 100644
--- a/pprobe/toggle/cli.py
+++ b/pprobe/toggle/cli.py
@@ -135,9 +135,7 @@ def _save_toggles_to_file(self, file_path, toggle_dict):
                     value_str = (
                         "true"
                         if value is True
-                        else "false"
-                        if value is False
-                        else value
+                        else "false" if value is False else value
                     )
                     file.write(f"{name}={value_str}\n")
         except Exception as e: