fix bug: OOM

Author: 545999961

FlagEmbedding/inference/embedder/decoder_only/base.py

@@ -257,7 +257,7 @@ def encode_single_device(
                 flag = True
             except RuntimeError as e:
                 batch_size = batch_size * 3 // 4
-            except torch.OutOfMemoryError as e:
+            except torch.cuda.OutOfMemoryError as e:
                 batch_size = batch_size * 3 // 4
 
         # encode

FlagEmbedding/inference/embedder/decoder_only/icl.py

@@ -409,7 +409,7 @@ def encode_queries_single_device(
                 flag = True
             except RuntimeError as e:
                 batch_size = batch_size * 3 // 4
-            except torch.OutOfMemoryError as e:
+            except torch.cuda.OutOfMemoryError as e:
                 batch_size = batch_size * 3 // 4
 
         # encode
@@ -519,7 +519,7 @@ def encode_single_device(
                 flag = True
             except RuntimeError as e:
                 batch_size = batch_size * 3 // 4
-            except torch.OutOfMemoryError as e:
+            except torch.cuda.OutOfMemoryError as e:
                 batch_size = batch_size * 3 // 4
 
         # encode

FlagEmbedding/inference/embedder/encoder_only/base.py

@@ -270,7 +270,7 @@ def encode_only(
                 flag = True
             except RuntimeError as e:
                 batch_size = batch_size * 3 // 4
-            except torch.OutOfMemoryError as e:
+            except torch.cuda.OutOfMemoryError as e:
                 batch_size = batch_size * 3 // 4
 
         # encode

FlagEmbedding/inference/embedder/encoder_only/m3.py

@@ -406,7 +406,7 @@ def _process_colbert_vecs(colbert_vecs: np.ndarray, attention_mask: list):
                 flag = True
             except RuntimeError as e:
                 batch_size = batch_size * 3 // 4
-            except torch.OutOfMemoryError as e:
+            except torch.cuda.OutOfMemoryError as e:
                 batch_size = batch_size * 3 // 4
 
         # encode

FlagEmbedding/inference/reranker/decoder_only/base.py

@@ -412,7 +412,7 @@ def compute_score_single_gpu(
                 flag = True
             except RuntimeError as e:
                 batch_size = batch_size * 3 // 4
-            except torch.OutOfMemoryError as e:
+            except torch.cuda.OutOfMemoryError as e:
                 batch_size = batch_size * 3 // 4
 
         dataset, dataloader = None, None

FlagEmbedding/inference/reranker/decoder_only/layerwise.py

@@ -282,7 +282,7 @@ def compute_score_single_gpu(
                 flag = True
             except RuntimeError as e:
                 batch_size = batch_size * 3 // 4
-            except torch.OutOfMemoryError as e:
+            except torch.cuda.OutOfMemoryError as e:
                 batch_size = batch_size * 3 // 4
 
         dataset, dataloader = None, None

FlagEmbedding/inference/reranker/decoder_only/lightweight.py

@@ -368,7 +368,7 @@ def compute_score_single_gpu(
                 flag = True
             except RuntimeError as e:
                 batch_size = batch_size * 3 // 4
-            except torch.OutOfMemoryError as e:
+            except torch.cuda.OutOfMemoryError as e:
                 batch_size = batch_size * 3 // 4
 
         all_scores = []

FlagEmbedding/inference/reranker/encoder_only/base.py

@@ -169,7 +169,7 @@ def compute_score_single_gpu(
                 flag = True
             except RuntimeError as e:
                 batch_size = batch_size * 3 // 4
-            except torch.OutOfMemoryError as e:
+            except torch.cuda.OutOfMemoryError as e:
                 batch_size = batch_size * 3 // 4
 
         all_scores = []