File tree Expand file tree Collapse file tree 1 file changed +3
-6
lines changed
torch/testing/_internal/distributed/_tensor Expand file tree Collapse file tree 1 file changed +3
-6
lines changed Original file line number Diff line number Diff line change 5858NUM_DEVICES = 4
5959
6060# We use this as a proxy for "multiple GPUs exist"
61- if (TEST_CUDA or TEST_XPU ) and DEVICE_COUNT > 1 :
61+ if (TEST_CUDA or TEST_XPU or TEST_HPU ) and DEVICE_COUNT > 1 :
6262 # when we actually have multiple GPUs, relax the requirement to smaller counts.
6363 NUM_DEVICES = min (NUM_DEVICES , DEVICE_COUNT )
6464
@@ -339,11 +339,8 @@ def world_size(self) -> int:
339339
340340 @property
341341 def device_type (self ) -> str :
342- # if enough GPU we can use GPU, otherwise we fallback to CPU
343- if (
344- not (TEST_CUDA or TEST_XPU )
345- or torch .accelerator .device_count () < self .world_size
346- ):
342+ # if enough GPU/XPU/HPU we can use those devices, otherwise we fallback to CPU
343+ if not (TEST_CUDA or TEST_XPU or TEST_HPU ) or DEVICE_COUNT < self .world_size :
347344 return "cpu"
348345 else :
349346 return DEVICE_TYPE
You can’t perform that action at this time.
0 commit comments