diff --git a/torchrec/distributed/comm_ops.py b/torchrec/distributed/comm_ops.py index 4d950c7e9..111970853 100644 --- 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 index 4c66511ef..64180c788 100644 --- 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 index ef6a67098..c2c1d6fce 100644 --- 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 index 0b554cd02..63792502d 100644 --- 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 index d09f30781..48de934df 100644 --- 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 index 94d7574f6..4296725bd 100644 --- 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 index 3e979b34d..628caf3e3 100644 --- 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 index ebdce6acb..59821b76b 100644 --- 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