magnet/eqV2/module_list.py

import torch


class ModuleListInfo(torch.nn.ModuleList):
    def __init__(self, info_str, modules=None):
        super().__init__(modules)
        self.info_str = str(info_str)
    

    def __repr__(self): 
        return self.info_str