def METHOD_NAME(modclass: Union[Type[nn.Module],
                                    Sequence[Type[nn.Module]]]):
    """
    Registers a PyTorch nn.Module (by class type) with a replacement that
    returns an LBANN Layer object (potentially corresponding to a subgraph).
    The decorator can be repeated to register multiple modules to the same
    replacement.
    Syntax::
        @register_module(ClassName)
        def replacement(mod: ClassName):
            return lbann.MatchingClass(...)
    :param modclass: A module type or a sequence thereof to be replaced.
    """
    def func(f):
        if isinstance(modclass, (list, tuple)):
            for cls in modclass:
                modules[cls] = f
        else:
            modules[modclass] = f
        return f
    return func