Fix ITEP tests in dev-nosan mode (pytorch#2242)

Summary: Pull Request resolved: pytorch#2242 Fixing ITEP tests in dev-nonsan mode Reviewed By: ge0405 Differential Revision: D59617484 fbshipit-source-id: 6f936e1bc208ff8c8650100b6abfa2c52caca452
TroyGarden · Jul 23, 2024 · 742acbb · 742acbb
1 parent 9264186
commit 742acbb
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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