From 6d599d53fd6d463e6cd933af7609ba3516414264 Mon Sep 17 00:00:00 2001 From: Laurent Date: Thu, 8 Feb 2024 13:59:07 +0000 Subject: [PATCH] beautify `EXPECTED_TREE` in test_chain.py --- tests/fluxion/layers/test_chain.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/tests/fluxion/layers/test_chain.py b/tests/fluxion/layers/test_chain.py index 03e861f..185be53 100644 --- a/tests/fluxion/layers/test_chain.py +++ b/tests/fluxion/layers/test_chain.py @@ -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: