transformer 意外终止

import math
import mindspore as ms
from mindspore import nn, ops


def get_mask(inputs):
    l_ = inputs.shape[1]
    mask = ops.ones((l_, l_), dtype=ms.float32) * (-1e+9)
    mask_flags = ops.sequence_mask(ops.arange(1, l_ + 1, dtype=ms.int32))
    mask[mask_flags] = 0.0
    return mask


def get_pad_mask(token):
    mask = ops.zeros_like(token, dtype=ms.float32)
    mask[token == 2] = -1e+9
    return mask


class PositionEncoder(nn.Cell):
    def __init__(self, d_model, max_len=1000):
        super().__init__()
        self.d_model = d_model
        self.max_len = max_len

        self.pe = ops.zeros((max_len, d_model), dtype=ms.float32)
        position = ops.arange(0, max_len, dtype=ms.float32).unsqueeze(-1)

        div_term = ops.exp(
            ops.arange(0, d_model, 2) * (-math.log(10000.0) / d_model)
        )
        self.pe[:, 0::2] = ops.sin(position * div_term)
        self.pe[:, 1::2] = ops.cos(position * div_term)

    def construct(self, inputs):
        return inputs + self.pe[:inputs.shape[0], :].unsqueeze(1)


class CopyModel(nn.Cell):
    def __init__(self, d_model, max_length):
        super().__init__()

        self.embed_src = nn.Embedding(max_length + 3, d_model)
        self.embed_dst = nn.Embedding(max_length + 3, d_model)
        self.positioner = PositionEncoder(d_model)

        self.transformer = nn.Transformer(d_model)

        # self.flatten = nn.Flatten()
        self.fc = nn.Dense(d_model, max_length + 3)

    def construct(self, inputs, targets):
        src = self.embed_src(inputs.transpose(1, 0))
        src = self.positioner(src)
        dst = self.embed_dst(targets.transpose(1, 0))
        dst = self.positioner(dst)

        dst_mask = get_mask(targets)
        src_pad_mask = get_pad_mask(inputs)

        dst_pad_mask = get_pad_mask(targets)

        out = self.transformer(src, dst, tgt_mask=dst_mask,
                               src_key_padding_mask=src_pad_mask, tgt_key_padding_mask=dst_pad_mask).transpose(1, 0, 2)

        out = self.fc(out).transpose(0, 2, 1)
        return out