def predict_fn()

in right_size_your_sagemaker_endpoints/inference.py [0:0]

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

def predict_fn(input_data, model):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model.to(device)
    model.eval()
    with torch.no_grad():
        return model(input_data.to(device))