IndexError while running the NMT example

[nashid]

I am running the NMT example but with a different dataset. During inference, I get the following error with top_k > 1:

site-packages/graph4nlp_cu111-0.4.0-py3.9.egg/graph4nlp/pytorch/models/graph2seq.py", line 251, in translate
return self.encoder_decoder_beam_search(
File "/xxxxxx/lib/python3.9/site-packages/graph4nlp_cu111-0.4.0-py3.9.egg/graph4nlp/pytorch/models/graph2seq.py", line in encoder_decoder_beam_search
beam_results = generator.generate(batch_graph=batch_graph, oov_dict=oov_dict, topk=topk)
File "/xxxxxx/lib/python3.9/site-packages/graph4nlp_cu111-0.4.0-py3.9.egg/graph4nlp/pytorch/modules/prediction/generation/decoder_strategy.py", line 150, in gener        ate
return self._beam_search(**params)
File "/xxxxxx/lib/python3.9/site-packages/graph4nlp_cu111-0.4.0-py3.9.egg/graph4nlp/pytorch/modules/prediction/generation/decoder_strategy.py", line 206, in _beam        _search
graph_node_mask[batch_idx, :].unsqueeze(0).expand(beam_size, -1).contiguous()
IndexError: too many indices for tensor of dimension 1

The problem is it does not happen for all datasets. But it is evident there is a subtle bug in the beam search. Can anyone please help?

[AlanSwift]

This issue will be closed. Feel free to reopen it if needed.