[ALL] Merge master (#447)

This commit is contained in:
Philippe Tillet
2022-01-30 20:21:20 -08:00
committed by GitHub
parent bef76b142a
commit 807d8a1945
14 changed files with 199 additions and 130 deletions

View File

@@ -25,13 +25,13 @@ def get_p2p_matrix():
def get_p2p_devices():
matrix = get_p2p_matrix()
idx = np.where(matrix == "OK")
return f"cuda:{idx[0][0]}", f"cuda:{idx[1][0]}"
return [f"cuda:{idx[0][0]}", f"cuda:{idx[1][0]}"] if len(idx[0]) > 0 else []
def get_non_p2p_devices():
matrix = get_p2p_matrix()
idx = np.where(matrix == "NS")
return f"cuda:{idx[0][0]}", f"cuda:{idx[1][0]}"
return [f"cuda:{idx[0][0]}", f"cuda:{idx[1][0]}"] if len(idx[0]) > 0 else []
p2p_devices = get_p2p_devices()