fix test_debug_print

Follow-up of #173
finegrain-ai · Jan 12, 2024 · dba9065 · dba9065
1 parent 457c3f5
commit dba9065
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/fluxion/layers/test_chain.py b/tests/fluxion/layers/test_chain.py
@@ -230,8 +230,8 @@ def test_setattr_dont_register() -> None:
 
 
 EXPECTED_TREE = (
-    "(CHAIN)\n    ├── Linear(in_features=1, out_features=1) (x2)\n    └── (CHAIN)\n        ├── Linear(in_features=1,"
-    " out_features=1) #1\n        └── Linear(in_features=2, out_features=1) #2"
+    "(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"
 )