summaryrefslogtreecommitdiff
path: root/test/test_cuda.py
diff options
context:
space:
mode:
authorSsnL <tongzhou.wang.1994@gmail.com>2018-12-09 21:10:39 -0800
committerFacebook Github Bot <facebook-github-bot@users.noreply.github.com>2018-12-09 21:12:46 -0800
commitfab80851111ffe7c090f05a051f49f1c54342815 (patch)
tree998271465c356aca6e9cb528c40621efc4cf7a87 /test/test_cuda.py
parent5fd69e75516326c2cdb3af9ffa702068adc6fbf6 (diff)
downloadpytorch-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.py2
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):