[MMS] Fix concatenation of emissions (#5133)

facebookresearch · May 23, 2023 · 87d3005 · 87d3005
1 parent af12c9c
commit 87d3005
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/examples/mms/data_prep/align_and_segment.py b/examples/mms/data_prep/align_and_segment.py
@@ -51,12 +51,12 @@ def generate_emissions(model, audio_file):
             offset = time_to_frame(input_start_time)
 
             emissions_ = emissions_[
-                :, emission_start_frame - offset : emission_end_frame - offset
+                emission_start_frame - offset : emission_end_frame - offset, :
             ]
             emissions_arr.append(emissions_)
             i += EMISSION_INTERVAL
 
-    emissions = torch.cat(emissions_arr, dim=1).squeeze()
+    emissions = torch.cat(emissions_arr, dim=0).squeeze()
     emissions = torch.log_softmax(emissions, dim=-1)
 
     stride = float(waveform.size(1) * 1000 / emissions.size(0) / SAMPLING_FREQ)