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 58
58
NUM_DEVICES = 4
59
59
60
60
# 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 :
62
62
# when we actually have multiple GPUs, relax the requirement to smaller counts.
63
63
NUM_DEVICES = min (NUM_DEVICES , DEVICE_COUNT )
64
64
@@ -339,11 +339,8 @@ def world_size(self) -> int:
339
339
340
340
@property
341
341
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 :
347
344
return "cpu"
348
345
else :
349
346
return DEVICE_TYPE
You can’t perform that action at this time.
0 commit comments