Skip to content

Commit

Permalink
[CUDA][P2P] Check device capability in requires_cuda_p2p_access (py…
Browse files Browse the repository at this point in the history
  • Loading branch information
eqy authored and pytorchmergebot committed Aug 30, 2024
1 parent 92f282c commit c07e566
Showing 1 changed file with 3 additions and 1 deletion.
4 changes: 3 additions & 1 deletion test/distributed/test_symmetric_memory.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,9 @@

def requires_cuda_p2p_access():
cuda_p2p_access_available = (
torch.cuda.is_available() and torch.cuda.device_count() >= 2
torch.cuda.is_available()
and torch.cuda.get_device_capability() >= (8, 0)
and torch.cuda.device_count() >= 2
)
num_devices = torch.cuda.device_count()
for i in range(num_devices - 1):
Expand Down

0 comments on commit c07e566

Please sign in to comment.