Fix

huggingface · Nov 13, 2023 · 5fedeeb · 5fedeeb
1 parent 93c7928
commit 5fedeeb
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/optimum/neuron/distributed/decoder_models.py b/optimum/neuron/distributed/decoder_models.py
@@ -247,7 +247,7 @@ def sequence_parallel_forward(
 
             # Reshape outputs
             if sequence_parallel_enabled:
-                # [batch, seq_len, num_attention_heads, head_size] -> [seq_len, batch, hidden_size]
+                # [batch, num_attention_heads, seq_len, head_size] -> [seq_len, batch, hidden_size]
                 attn_output = attn_output.permute(2, 0, 1, 3).contiguous()
                 attn_output = attn_output.view(*attn_output.shape[:2], -1)
             else:

diff --git a/optimum/neuron/utils/cache_utils.py b/optimum/neuron/utils/cache_utils.py
@@ -405,7 +405,7 @@ def add_in_registry(repo_id: str, neuron_hash: "NeuronHash"):
                     commit_message=f"Add {model_name_or_path} in registry for NeuronHash {overall_hash}",
                     parent_commit=head,
                 )
-            except ValueError as e:
+            except Exception as e:
                 if "A commit has happened since" in str(e):
                     logger.info(
                         "A commit has happened in cache repository since we tried to update the registry, starting again..."