core.models.equiformer_v2.module_list

Classes

ModuleListInfo

Holds submodules in a list.

Module Contents

class core.models.equiformer_v2.module_list.ModuleListInfo(info_str, modules=None)

Bases: torch.nn.ModuleList

Holds submodules in a list.

ModuleList can be indexed like a regular Python list, but modules it contains are properly registered, and will be visible by all Module methods.

Parameters:

modules (iterable, optional) – an iterable of modules to add

Example:

class MyModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.linears = nn.ModuleList([nn.Linear(10, 10) for i in range(10)])

    def forward(self, x):
        # ModuleList can act as an iterable, or be indexed using ints
        for i, l in enumerate(self.linears):
            x = self.linears[i // 2](x) + l(x)
        return x

info_str

__repr__() → str

Return a custom repr for ModuleList that compresses repeated module representations.