Update align_and_segment.py (#5317)

Fix MMS alignment code
facebookresearch · Sep 7, 2023 · b5d89cd · b5d89cd
1 parent 4db2649
commit b5d89cd
Showing 1 changed file with 6 additions and 5 deletions.
diff --git a/examples/mms/data_prep/align_and_segment.py b/examples/mms/data_prep/align_and_segment.py
@@ -87,13 +87,14 @@ def get_alignments(
     blank = dictionary["<blank>"]
 
     targets = torch.tensor(token_indices, dtype=torch.int32).to(DEVICE)
-    input_lengths = torch.tensor(emissions.shape[0])
-    target_lengths = torch.tensor(targets.shape[0])
-
+
+    input_lengths = torch.tensor(emissions.shape[0]).unsqueeze(-1)
+    target_lengths = torch.tensor(targets.shape[0]).unsqueeze(-1)
     path, _ = F.forced_align(
-        emissions, targets, input_lengths, target_lengths, blank=blank
+        emissions.unsqueeze(0), targets.unsqueeze(0), input_lengths, target_lengths, blank=blank
     )
-    path = path.to("cpu").tolist()
+    path = path.squeeze().to("cpu").tolist()
+
     segments = merge_repeats(path, {v: k for k, v in dictionary.items()})
     return segments, stride