fix OOM bug: direct exception

Author: hanhainebula

FlagEmbedding/inference/embedder/decoder_only/base.py

@@ -180,9 +180,7 @@ def encode_single_device(
                 last_hidden_state = self.model(**test_inputs_batch, return_dict=True).last_hidden_state
                 embeddings = last_token_pool(last_hidden_state, test_inputs_batch['attention_mask'])
                 flag = True
-            except RuntimeError as e:
-                batch_size = batch_size * 3 // 4
-            except torch.OutOfMemoryError as e:
+            except:
                 batch_size = batch_size * 3 // 4
 
         # encode

FlagEmbedding/inference/embedder/decoder_only/icl.py

@@ -278,9 +278,7 @@ def encode_queries_single_device(
                 last_hidden_state = self.model(**test_inputs_batch, return_dict=True).last_hidden_state
                 embeddings = last_token_pool(last_hidden_state, test_inputs_batch['attention_mask'])
                 flag = True
-            except RuntimeError as e:
-                batch_size = batch_size * 3 // 4
-            except torch.OutOfMemoryError as e:
+            except:
                 batch_size = batch_size * 3 // 4
 
         # encode
@@ -391,9 +389,7 @@ def encode_single_device(
                 last_hidden_state = self.model(**test_inputs_batch, return_dict=True).last_hidden_state
                 embeddings = last_token_pool(last_hidden_state, test_inputs_batch['attention_mask'])
                 flag = True
-            except RuntimeError as e:
-                batch_size = batch_size * 3 // 4
-            except torch.OutOfMemoryError as e:
+            except:
                 batch_size = batch_size * 3 // 4
 
         # encode

FlagEmbedding/inference/embedder/encoder_only/base.py

@@ -170,9 +170,7 @@ def encode_single_device(
                 last_hidden_state = self.model(**test_inputs_batch, return_dict=True).last_hidden_state
                 embeddings = self.pooling(last_hidden_state, test_inputs_batch['attention_mask'])
                 flag = True
-            except RuntimeError as e:
-                batch_size = batch_size * 3 // 4
-            except torch.OutOfMemoryError as e:
+            except:
                 batch_size = batch_size * 3 // 4
 
         # encode

FlagEmbedding/inference/embedder/encoder_only/m3.py

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

FlagEmbedding/inference/reranker/decoder_only/base.py

@@ -317,9 +317,7 @@ def compute_score_single_gpu(
 
                 self.model(**batch_inputs, output_hidden_states=True)
                 flag = True
-            except RuntimeError as e:
-                batch_size = batch_size * 3 // 4
-            except torch.OutOfMemoryError as e:
+            except:
                 batch_size = batch_size * 3 // 4
 
         dataset, dataloader = None, None

FlagEmbedding/inference/reranker/decoder_only/layerwise.py

@@ -225,9 +225,7 @@ def compute_score_single_gpu(
 
                 self.model(**batch_inputs, output_hidden_states=True, cutoff_layers=cutoff_layers)
                 flag = True
-            except RuntimeError as e:
-                batch_size = batch_size * 3 // 4
-            except torch.OutOfMemoryError as e:
+            except:
                 batch_size = batch_size * 3 // 4
 
         dataset, dataloader = None, None

FlagEmbedding/inference/reranker/decoder_only/lightweight.py

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

FlagEmbedding/inference/reranker/encoder_only/base.py

@@ -134,9 +134,7 @@ def compute_score_single_gpu(
                 ).to(device)
                 scores = self.model(**test_inputs_batch, return_dict=True).logits.view(-1, ).float()
                 flag = True
-            except RuntimeError as e:
-                batch_size = batch_size * 3 // 4
-            except torch.OutOfMemoryError as e:
+            except:
                 batch_size = batch_size * 3 // 4
 
         all_scores = []