(self)
| 181 | self.assertEqual(registry._hook_order, ["multiply_hook"]) |
| 182 | |
| 183 | def test_stateful_hook(self): |
| 184 | registry = HookRegistry.check_if_exists_or_initialize(self.model) |
| 185 | registry.register_hook(StatefulAddHook(1), "stateful_add_hook") |
| 186 | |
| 187 | self.assertEqual(registry.hooks["stateful_add_hook"].increment, 0) |
| 188 | |
| 189 | input = torch.randn(1, 4, device=torch_device, generator=self.get_generator()) |
| 190 | num_repeats = 3 |
| 191 | |
| 192 | for i in range(num_repeats): |
| 193 | result = self.model(input) |
| 194 | if i == 0: |
| 195 | output1 = result |
| 196 | |
| 197 | self.assertEqual(registry.get_hook("stateful_add_hook").increment, num_repeats) |
| 198 | |
| 199 | registry.reset_stateful_hooks() |
| 200 | output2 = self.model(input) |
| 201 | |
| 202 | self.assertEqual(registry.get_hook("stateful_add_hook").increment, 1) |
| 203 | self.assertTrue(torch.allclose(output1, output2)) |
| 204 | |
| 205 | def test_inference(self): |
| 206 | registry = HookRegistry.check_if_exists_or_initialize(self.model) |
nothing calls this directly
no test coverage detected