247203f29b7e841204c76d922c1ea5b2680c3663,reagent/model_utils/seq2slate_utils.py,,per_symbol_to_per_seq_probs,#Any#Any#,115

Before Change

    // tgt_out_idx shape: batch_size, seq_len
    // output shape: batch_size, 1
    return torch.prod(
        torch.gather(per_symbol_probs, , tgt_out_idx.unsqueeze(-1)).squeeze(2),
        dim=1,
        keepdim=True,
    )

After Change

    // output shape: batch_size, 1
    return (
        torch.prod(
            torch.gather(per_symbol_probs, 2, tgt_out_idx.unsqueeze(2)).squeeze(2),
            dim=1,
            keepdim=True,
        )

Instances