From 86137432ca5dc4834c013255713f9daa67675b3e Mon Sep 17 00:00:00 2001 From: Xiaodong Wang Date: Wed, 19 Jun 2024 01:38:44 -0700 Subject: [PATCH] Unify the cuda and hip path for fbgemm in trainer embedding pruning Differential Revision: D58775111 --- torchrec/modules/itep_modules.py | 11 +++-------- 1 file changed, 3 insertions(+), 8 deletions(-) diff --git a/torchrec/modules/itep_modules.py b/torchrec/modules/itep_modules.py index 945d08eef..5a477fc2e 100644 --- a/torchrec/modules/itep_modules.py +++ b/torchrec/modules/itep_modules.py @@ -17,14 +17,9 @@ from torchrec.sparse.jagged_tensor import _pin_and_move, _to_offsets, KeyedJaggedTensor try: - if torch.version.hip: - torch.ops.load_library( - "//deeplearning/fbgemm/fbgemm_gpu:intraining_embedding_pruning_hip" - ) - else: - torch.ops.load_library( - "//deeplearning/fbgemm/fbgemm_gpu:intraining_embedding_pruning_cuda" - ) + torch.ops.load_library( + "//deeplearning/fbgemm/fbgemm_gpu:intraining_embedding_pruning" + ) except OSError: pass