fix issue

clemente0731 · Jul 4, 2024 · 8e350c3 · 8e350c3
1 parent b763664
commit 8e350c3
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 7 deletions.
diff --git 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
@@ -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: