fix rebase and CI

Signed-off-by: Yu Zhentao <[email protected]>
huggingface · Dec 4, 2024 · 6f12c92 · 6f12c92
1 parent 1a70b12
commit 6f12c92
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/optimum/habana/transformers/generation/utils.py b/optimum/habana/transformers/generation/utils.py
@@ -120,6 +120,7 @@
 # Initial generated token index is set to 1 to accomodate SOS (start of string) token.
 INITIAL_TOKEN_IDX = 1
 
+logging.set_verbosity_info()
 logger = logging.get_logger(__name__)
 
 

diff --git a/optimum/habana/transformers/models/modeling_all_models.py b/optimum/habana/transformers/models/modeling_all_models.py
@@ -55,6 +55,9 @@ def allocate(self, inp_seq_len, dtype, device, shape):
 
     @staticmethod
     def update(prev, cur, dim, idx, inp_seq_len):
+        cur = cur.to(prev.device)
+        if idx is not None:
+            idx = idx.to(prev.device)
         orig_cur = cur
         if prev.shape == cur.shape:
             prev.copy_(cur)