diff --git a/torchrec/modules/itep_modules.py b/torchrec/modules/itep_modules.py
index 14a479b2e..a3c428834 100644
--- a/torchrec/modules/itep_modules.py
+++ b/torchrec/modules/itep_modules.py
@@ -19,7 +19,7 @@
 
 try:
     torch.ops.load_library(
-        "//deeplearning/fbgemm/fbgemm_gpu:intraining_embedding_pruning"
+        "//deeplearning/fbgemm/fbgemm_gpu:intraining_embedding_pruning_gpu"
     )
 except OSError:
     pass