how to register custom hook in mmpretrain to save model embeddings

[sminut]

I'm trying to save the EfficientNet model embeddings to file. I followed the "Custom Hooks" section in this tutorial: https://mmengine.readthedocs.io/en/latest/tutorials/hook.html and I created a file called "save_embeddings_hook.py" with the following hook:

import torch
from mmengine.registry import HOOKS
from mmengine.hooks import Hook

@HOOKS.register_module()
class SaveEmbeddingsHook(Hook):
    def __init__(self, save_interval, out_dir, priority, stage):
        self.save_interval = save_interval
        self.out_dir = out_dir
        self.priority = priority
        self.stage = stage

    def after_test_iter(self, runner, batch_idx, data_batch, outputs):
        if self.stage == 'test':
            embeddings = runner.model.backbone.features(data_batch['inputs'])
            for i, data_sample in enumerate(data_batch['data_samples']):
                img_path = data_sample.get('img_path', f'batch_{batch_idx}_img_{i}')
                save_path = f'{self.out_dir}/embedding_{img_path.stem}.pt'
                torch.save(embeddings[i].cpu(), save_path)

Next, in config.py I put

custom_hooks = [
    dict(
        type='SaveEmbeddingsHook',
        save_interval=1,
        out_dir='embeddings',
        priority='ABOVE_NORMAL',
        stage='test'
    )
]

Then, in the docker container I run test.py as follows:

python3 /workspace/mmpretrain/tools/test.py config.py output/epoch_50.pth

but I'm getting the following error:

Traceback (most recent call last):
  File "/workspace/mmpretrain/tools/test.py", line 193, in <module>
    main()
  File "/workspace/mmpretrain/tools/test.py", line 175, in main
    runner = Runner.from_cfg(cfg)
  File "/opt/conda/lib/python3.10/site-packages/mmengine/runner/runner.py", line 462, in from_cfg
    runner = cls(
  File "/opt/conda/lib/python3.10/site-packages/mmengine/runner/runner.py", line 442, in __init__
    self.register_hooks(default_hooks, custom_hooks)
  File "/opt/conda/lib/python3.10/site-packages/mmengine/runner/runner.py", line 1995, in register_hooks
    self.register_custom_hooks(custom_hooks)
  File "/opt/conda/lib/python3.10/site-packages/mmengine/runner/runner.py", line 1976, in register_custom_hooks
    self.register_hook(hook)
  File "/opt/conda/lib/python3.10/site-packages/mmengine/runner/runner.py", line 1877, in register_hook
    hook_obj = HOOKS.build(hook)
  File "/opt/conda/lib/python3.10/site-packages/mmengine/registry/registry.py", line 570, in build
    return self.build_func(cfg, *args, **kwargs, registry=self)
  File "/opt/conda/lib/python3.10/site-packages/mmengine/registry/build_functions.py", line 102, in build_from_cfg
    raise KeyError(
KeyError: 'SaveEmbeddingsHook is not in the mmpretrain::hook registry. Please check whether the value of `SaveEmbeddingsHook` is correct or it was registered as expected. More details can be found at https://mmengine.readthedocs.io/en/latest/advanced_tutorials/config.html#import-the-custom-module'

So it seems it can't find my hook. The file I saved it in is in the same directory as config.py.

How do I properly register the custom hook?

Thanks!