diff options
author | SsnL <tongzhou.wang.1994@gmail.com> | 2018-12-09 21:10:39 -0800 |
---|---|---|
committer | Facebook Github Bot <facebook-github-bot@users.noreply.github.com> | 2018-12-09 21:12:46 -0800 |
commit | fab80851111ffe7c090f05a051f49f1c54342815 (patch) | |
tree | 998271465c356aca6e9cb528c40621efc4cf7a87 /test/test_cuda.py | |
parent | 5fd69e75516326c2cdb3af9ffa702068adc6fbf6 (diff) | |
download | pytorch-fab80851111ffe7c090f05a051f49f1c54342815.tar.gz pytorch-fab80851111ffe7c090f05a051f49f1c54342815.tar.bz2 pytorch-fab80851111ffe7c090f05a051f49f1c54342815.zip |
_get_device_index supports parsing device strings
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/14929
Reviewed By: weiyangfb
Differential Revision: D13394498
Pulled By: soumith
fbshipit-source-id: 948c6118abdf6c1e1a8a17709333954cafb2345e
Diffstat (limited to 'test/test_cuda.py')
-rw-r--r-- | test/test_cuda.py | 2 |
1 files changed, 1 insertions, 1 deletions
diff --git a/test/test_cuda.py b/test/test_cuda.py index 4867ed6188..8226479a67 100644 --- a/test/test_cuda.py +++ b/test/test_cuda.py @@ -787,7 +787,7 @@ class TestCuda(TestCase): # interlace torch.cuda.empty_cache() - gen0 = self._test_memory_stats_generator(self, device=0, N=35) + gen0 = self._test_memory_stats_generator(self, device='cuda:0', N=35) gen1 = self._test_memory_stats_generator(self, device=torch.device('cuda:1'), N=35) end0 = end1 = False while not (end0 and end1): |