From 5fedeeb69d2ba80942a3b943d365c41fdaff25bd Mon Sep 17 00:00:00 2001 From: Michael Benayoun Date: Mon, 13 Nov 2023 18:52:23 +0100 Subject: [PATCH] Fix --- optimum/neuron/distributed/decoder_models.py | 2 +- optimum/neuron/utils/cache_utils.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/optimum/neuron/distributed/decoder_models.py b/optimum/neuron/distributed/decoder_models.py index 11a8179b6..bad00254b 100644 --- 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 index 45db6c7ea..e42c897a0 100644 --- 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..."