models/connector.py

from torch import nn




def connector(connector_type='linear', **kwargs):
    print("Build connector:", connector_type)
    if connector_type == 'linear':
        return nn.ModuleList([nn.Linear(kwargs['input_dim'], kwargs['output_dim']) for i in range(kwargs['num_layers'])])
    else:
        raise NotImplemented