Pytorch nn.TransformerEncoder «src_key_padding_mask» не работает должным образом

Я работаю с модулем Pytorch nn.TransformerEncoder. Я получил входные образцы с (как обычно) формой (batch-size, seq-len, emb-dim). Все образцы в одном пакете были дополнены нулями до размера самого большого образца в этом пакете. Поэтому я хочу, чтобы внимание всех нулевых значений игнорировалось.

В документации сказано, что нужно добавить аргумент src_key_padding_mask к функции forward модуля nn.TransformerEncoder. Эта маска должна быть тензором формы (batch-size, seq-len) и иметь для каждого индекса либо True для нулей, либо False для чего-то еще.

Я добился этого, выполнив:

. . .

def forward(self, x):
    # x.size -> i.e.: (200, 28, 200)

    mask = (x == 0).cuda().reshape(x.shape[0], x.shape[1])
    # mask.size -> i.e.: (200, 20)

    x = self.embed(x.type(torch.LongTensor).to(device=device))
    x = self.pe(x)

    x = self.transformer_encoder(x, src_key_padding_mask=mask)

    . . .

Все работает хорошо, когда я не устанавливаю src_key_padding_mask. Но ошибка, которую я получаю, когда я это делаю, заключается в следующем:

File "/home/me/.conda/envs/py37/lib/python3.7/site-packages/torch/nn/functional.py", line 4282, in multi_head_attention_forward
    assert key_padding_mask.size(0) == bsz
AssertionError

Похоже, что он сравнивает первое измерение маски, которое представляет собой размер партии, с bsz, который, вероятно, означает размер партии. Но почему тогда не получается? Помощь очень ценится!