def load_model()

in ai-judge/inference.py [0:0]

• 8 lines of code
• 2 McCabe index (conditional complexity)

def load_model(gpu=False):
    '''Load pretrained pytorch model'''
    model = CNN()
    model = model.double()
    model.load_state_dict(torch.load('pretrained_model'))
    if gpu:
        model = model.cuda()
        
    return model