from networks.mynet import TwoBranch
from networks_time.mynet import DiffTwoBranch


def get_network(args, ):
    if args.use_time_model:
        network = DiffTwoBranch(args)
    else:
        network = TwoBranch(args)
    

    n_parameters = sum(p.numel() for p in network.parameters() if p.requires_grad)
    print('number of params: %.2f M' % (n_parameters / 1024 / 1024))

    return network