nit

a-r-r-o-w · a-r-r-o-w · commit 7452ba5b2cc5 · 2025-08-06T11:51:56.000+02:00
diff --git a/tests/hooks/test_group_offloading.py b/tests/hooks/test_group_offloading.py
@@ -348,8 +348,6 @@ def apply_layer_output_tracker_hook(model: DummyModelWithLayerNorm):
 
         for (ref_name, ref_module), (name, module) in zip(model_ref.named_modules(), model.named_modules()):
             assert ref_name == name
-            if not isinstance(ref_module, (torch.nn.Linear, torch.nn.LayerNorm)):
-                continue
             ref_outputs = (
                 HookRegistry.check_if_exists_or_initialize(ref_module).get_hook("layer_output_tracker").outputs
             )

Original file line number	Diff line number	Diff line change
`@@ -348,8 +348,6 @@ def apply_layer_output_tracker_hook(model: DummyModelWithLayerNorm):`
`348`	`348`
`349`	`349`	`for (ref_name, ref_module), (name, module) in zip(model_ref.named_modules(), model.named_modules()):`
`350`	`350`	`assert ref_name == name`
`351`		`- if not isinstance(ref_module, (torch.nn.Linear, torch.nn.LayerNorm)):`
`352`		`- continue`
`353`	`351`	`ref_outputs = (`
`354`	`352`	`HookRegistry.check_if_exists_or_initialize(ref_module).get_hook("layer_output_tracker").outputs`
`355`	`353`	`)`