def prepare_model()

in 01-byoc/code/workflow.py [0:0]

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

def prepare_model(params):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(params)
    print("build model...")
    model = None
    if params.resume:
        model = torch.load(params.resume)
    elif params.model_name == 'VGGish':
        model = VGGish(params)

    model = model.to(device)
    return model