Skip to content

Load EI model to CPU by default in model_fn #68

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Apr 14, 2020
Merged

Load EI model to CPU by default in model_fn #68

Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
e16f49b
Load EI model to CPU by default in model_fn
dfan Apr 12, 2020
6d09a6b
Merge branch 'master' into pytorch_ei
chuyang-deng Apr 13, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion src/sagemaker_pytorch_serving_container/default_inference_handler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,8 @@ def default_model_fn(self, model_dir):
if not os.path.exists(model_path):
raise FileNotFoundError("Failed to load model with default model_fn: missing file {}."
.format(DEFAULT_MODEL_FILENAME))
return torch.jit.load(model_path)
# Client-framework is CPU only. But model will run in Elastic Inference server with CUDA.
return torch.jit.load(model_path, map_location=torch.device('cpu'))
else:
raise NotImplementedError(textwrap.dedent("""
Please provide a model_fn implementation.
Expand Down