Fix exception handling for torch.ops.load_libraries

nipung90 · facebook-github-bot · commit 52dc11aa22e6 · 2025-07-18T09:07:29.000-07:00
Differential Revision: D78563082
diff --git a/torchrec/distributed/comm_ops.py b/torchrec/distributed/comm_ops.py
@@ -25,7 +25,7 @@
 try:
     torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops")
     torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops_cpu")
-except OSError:
+except (OSError, RuntimeError):
     pass
 
 
diff --git a/torchrec/distributed/dist_data.py b/torchrec/distributed/dist_data.py
@@ -38,7 +38,7 @@
     torch.ops.load_library(
         "//deeplearning/fbgemm/fbgemm_gpu:merge_pooled_embeddings_cpu"
     )
-except OSError:
+except (OSError, RuntimeError):
     pass
 
 # OSS
diff --git a/torchrec/distributed/embedding.py b/torchrec/distributed/embedding.py
@@ -105,7 +105,7 @@
 try:
     torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops")
     torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops_cpu")
-except OSError:
+except (OSError, RuntimeError):
     pass
 
 
diff --git a/torchrec/distributed/embeddingbag.py b/torchrec/distributed/embeddingbag.py
@@ -113,7 +113,7 @@
     torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops")
     torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops_cpu")
     torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu/codegen:index_select_ops")
-except OSError:
+except (OSError, RuntimeError):
     pass
 
 
diff --git a/torchrec/distributed/model_parallel.py b/torchrec/distributed/model_parallel.py
@@ -58,7 +58,7 @@
 try:
     torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops")
     torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops_cpu")
-except OSError:
+except (OSError, RuntimeError):
     pass
 
 
diff --git a/torchrec/distributed/quant_embedding.py b/torchrec/distributed/quant_embedding.py
@@ -100,7 +100,7 @@
 try:
     torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops")
     torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops_cpu")
-except OSError:
+except (OSError, RuntimeError):
     pass
 
 
diff --git a/torchrec/quant/embedding_modules.py b/torchrec/quant/embedding_modules.py
@@ -79,7 +79,7 @@
 try:
     torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops")
     torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops_cpu")
-except OSError:
+except (OSError, RuntimeError):
     pass
 
 # OSS
diff --git a/torchrec/sparse/jagged_tensor.py b/torchrec/sparse/jagged_tensor.py
@@ -44,7 +44,7 @@
     torch.ops.load_library(
         "//deeplearning/fbgemm/fbgemm_gpu:permute_multi_embedding_ops_gpu"
     )
-except OSError:
+except (OSError, RuntimeError):
     pass
 
 

Original file line number	Diff line number	Diff line change
`@@ -38,7 +38,7 @@`
`38`	`38`	`torch.ops.load_library(`
`39`	`39`	`"//deeplearning/fbgemm/fbgemm_gpu:merge_pooled_embeddings_cpu"`
`40`	`40`	`)`
`41`		`-except OSError:`
	`41`	`+except (OSError, RuntimeError):`
`42`	`42`	`pass`
`43`	`43`
`44`	`44`	`# OSS`
Original file line number	Diff line number	Diff line change
`@@ -44,7 +44,7 @@`
`44`	`44`	`torch.ops.load_library(`
`45`	`45`	`"//deeplearning/fbgemm/fbgemm_gpu:permute_multi_embedding_ops_gpu"`
`46`	`46`	`)`
`47`		`-except OSError:`
	`47`	`+except (OSError, RuntimeError):`
`48`	`48`	`pass`
`49`	`49`
`50`	`50`