def model_fn(model_dir):
    'Load the PyTorch model from the `model_dir` directory.'
    print('Loading model.')
    model_info = {}
    model_info_path = os.path.join(model_dir, 'model_info.pth')
    with open(model_info_path, 'rb') as f:
        model_info = torch.load(f)
    print('model_info: {}'.format(model_info))
    device = torch.device(('cuda' if torch.cuda.is_available() else 'cpu'))
    model = BinaryClassifier(model_info['input_features'], model_info['hidden_dim'], model_info['output_dim'])
    model_path = os.path.join(model_dir, 'model.pth')
    with open(model_path, 'rb') as f:
        model.load_state_dict(torch.load(f))
    model.to(device).eval()
    print('Done loading model.')
    return model
Load the PyTorch model from the `model_dir` directory.
def model_fn(model_dir):
    
    print('Loading model.')
    model_info = {}
    model_info_path = os.path.join(model_dir, 'model_info.pth')
    with open(model_info_path, 'rb') as f:
        model_info = torch.load(f)
    print('model_info: {}'.format(model_info))
    device = torch.device(('cuda' if torch.cuda.is_available() else 'cpu'))
    model = BinaryClassifier(model_info['input_features'], model_info['hidden_dim'], model_info['output_dim'])
    model_path = os.path.join(model_dir, 'model.pth')
    with open(model_path, 'rb') as f:
        model.load_state_dict(torch.load(f))
    model.to(device).eval()
    print('Done loading model.')
    return model