From 928da1ee1c9ec964ec945ae634b6e36a1da69aab Mon Sep 17 00:00:00 2001 From: Pierre Chapuis Date: Thu, 8 Aug 2024 15:26:55 +0200 Subject: [PATCH] fix repr when there are non-Fluxion modules in a tree --- src/refiners/fluxion/layers/module.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/refiners/fluxion/layers/module.py b/src/refiners/fluxion/layers/module.py index 46b31cd..ba5287f 100644 --- a/src/refiners/fluxion/layers/module.py +++ b/src/refiners/fluxion/layers/module.py @@ -358,7 +358,10 @@ class ModuleTree: node: TreeNode = {"value": value, "class_name": class_name, "children": []} for child in module.children(): - node["children"].append(self._module_to_tree(module=child)) # type: ignore + if isinstance(child, Module): + node["children"].append(self._module_to_tree(module=child)) + else: + node["children"].append({"value": str(child), "class_name": child.__class__.__name__, "children": []}) return node def _fold_successive_identical(self, node: TreeNode) -> None: