def get_network()

in jcm/models/ncsnv2.py [0:0]

• 11 lines of code
• 4 McCabe index (conditional complexity)

def get_network(config):
    if config.data.image_size < 96:
        return functools.partial(NCSNv2, config=config)
    elif 96 <= config.data.image_size <= 128:
        return functools.partial(NCSNv2_128, config=config)
    elif 128 < config.data.image_size <= 256:
        return functools.partial(NCSNv2_256, config=config)
    else:
        raise NotImplementedError(
            f"No network suitable for {config.data.image_size}px implemented yet."
        )