diff options
| author | Lu Fang <lufang@fb.com> | 2018-12-03 14:07:50 -0800 |
|---|---|---|
| committer | Facebook Github Bot <facebook-github-bot@users.noreply.github.com> | 2018-12-03 14:10:30 -0800 |
| commit | e0f68671bdadee48670931d2b95188197ee6856a (patch) | |
| tree | d7c7518ab329d270621ef10df3cdeda875c67257 /caffe2/proto | |
| parent | b8da44dc13f41b7de1ff8db1b4ca8c4319004947 (diff) | |
| download | pytorch-e0f68671bdadee48670931d2b95188197ee6856a.tar.gz pytorch-e0f68671bdadee48670931d2b95188197ee6856a.tar.bz2 pytorch-e0f68671bdadee48670931d2b95188197ee6856a.zip | |
Restore device when import jit script module (#14454)
Summary:
We align the restore logic to `torch.load`, we try to restore to the right device, and if the device is not available, an exception is raised. We allow user to remap the device through a parameter `map_location`, it can be 1) a string like 'cuda:0`, `cpu`, 2) a device, torch.device('cpu'), 3) a dict, {'cuda:1', 'cuda:0'}, and a function, and its signature looks like string map_location(tensor, saved_device_string).
Pull Request resolved: https://github.com/pytorch/pytorch/pull/14454
Reviewed By: zrphercule
Differential Revision: D13271956
Pulled By: houseroad
fbshipit-source-id: dfd6b6049b0dc07549ddeddf2dea03ac53ba6d49
Diffstat (limited to 'caffe2/proto')
| -rw-r--r-- | caffe2/proto/torch.proto | 4 |
1 files changed, 3 insertions, 1 deletions
diff --git a/caffe2/proto/torch.proto b/caffe2/proto/torch.proto index 78e4c0d7bc..0dbca7cdd0 100644 --- a/caffe2/proto/torch.proto +++ b/caffe2/proto/torch.proto @@ -18,7 +18,9 @@ message TensorDef { optional RecordRef data = 6; - // future: device options + // device field stores the canonical device string, and it follows the + // format below: `(cpu|cuda)[:<device-index>]`, e.g., 'cuda:0' + optional string device = 7; } message ParameterDef { |