fix bugs of Outofmemory error

FlagOpen · Nov 29, 2024 · 9544ed6 · 9544ed6
1 parent 777a9ac
commit 9544ed6
Show file tree

Hide file tree

Showing 8 changed files with 9 additions and 9 deletions.
diff --git a/FlagEmbedding/inference/embedder/decoder_only/base.py b/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

diff --git a/FlagEmbedding/inference/embedder/decoder_only/icl.py b/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

diff --git a/FlagEmbedding/inference/embedder/encoder_only/base.py b/FlagEmbedding/inference/embedder/encoder_only/base.py
@@ -238,7 +238,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

diff --git a/FlagEmbedding/inference/embedder/encoder_only/m3.py b/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

diff --git a/FlagEmbedding/inference/reranker/decoder_only/base.py b/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

diff --git a/FlagEmbedding/inference/reranker/decoder_only/layerwise.py b/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

diff --git a/FlagEmbedding/inference/reranker/decoder_only/lightweight.py b/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 = []

diff --git a/FlagEmbedding/inference/reranker/encoder_only/base.py b/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 = []