use broadcasting in viterbi
@@ -222,9 +222,8 @@ def viterbi_decode(tag_sequence: torch.Tensor, transition_matrix: torch.Tensor):
# Evaluate the scores for all possible paths.
for timestep in range(1, sequence_length):
- # TODO(Mark): Use broadcasting here once Pytorch 0.2 is released.
- tiled_path_scores = path_scores[timestep - 1].expand_as(transition_matrix).transpose(0, 1)
- summed_potentials = tiled_path_scores + transition_matrix
+ # Add pairwise potentials to current scores.
+ summed_potentials = path_scores[timestep - 1].unsqueeze(-1) + transition_matrix
scores, paths = torch.max(summed_potentials, 0)
path_scores.append(tag_sequence[timestep, :] + scores.squeeze())
path_indices.append(paths.squeeze())
