beautify EXPECTED_TREE in test_chain.py

This commit is contained in:
Laurent 2024-02-08 13:59:07 +00:00 committed by Laureηt
parent e36dda63fd
commit 6d599d53fd

View file

@ -291,10 +291,13 @@ def test_setattr_dont_register() -> None:
assert module_keys(chain=chain) == ["Linear_1", "Linear_2"]
EXPECTED_TREE = (
"(CHAIN)\n ├── Linear(in_features=1, out_features=1, device=cpu, dtype=float32) (x2)\n └── (CHAIN)\n ├── Linear(in_features=1,"
" out_features=1, device=cpu, dtype=float32) #1\n └── Linear(in_features=2, out_features=1, device=cpu, dtype=float32) #2"
)
EXPECTED_TREE = """
(CHAIN)
Linear(in_features=1, out_features=1, device=cpu, dtype=float32) (x2)
(CHAIN)
Linear(in_features=1, out_features=1, device=cpu, dtype=float32) #1
Linear(in_features=2, out_features=1, device=cpu, dtype=float32) #2
""".strip()
def test_debug_print() -> None: