Оптимизация нейронных сеток под GPU, что будет узким местом?

Обучаю нейронную сетку, типа трансформер сделанную на Python, пока на CPU i5 12gen. У этого процессора 16 потоков, и по монитору видно, что они все суммарно загружаются на ~60%. Вопрос, будет ли, и существенно ли, увеличение скорости при использовании графической карты например RTX 4090 или даже нескольких?

У меня сомнения возникают, поскольку я вижу, что на CPU он не может использовать все ресурсы и возможно что-то является узким местом кроме как мультиядерности. Процесс существенно простой, кроме расчета сетки практически ничего нет.

От величины сетки процент загрузки практически не меняется. Сейчас пробую сетку на ~200млн параметров. Памяти пока хватает, потребляет максимум 30-40гб из 64гб в моменты отправки пачки на обучение.

Вот так выглядит код класса сетки, учится через MSELoss && Adam. Про загрузку в 60% речь идет когда выполняет код сетки. В данном случае около 50млн. параметров. На мой взгляд начинающего сеточника, здесь кажется особо оптимизировать нечего:

class TransformerModel(nn.Module):

    def __init__(self, in_vocab_size, sz_in, sz_out = 1):
        super(TransformerModel, self).__init__()

        emb_sz = 512
        n_heads = 16
        c_layers = 32
        c_feedForward = 512
        dropout = 0.0

        self.d_model = emb_sz
        self.sz_in = sz_in
        self.sz_out = sz_out

        self.embedding = nn.Embedding(
            num_embeddings=in_vocab_size,
            embedding_dim=emb_sz)

        self.pos_encoder = PositionalEncoding(
            d_model=emb_sz,
            dropout=dropout)

        self.transformer_layer = nn.TransformerEncoderLayer(
            d_model=emb_sz,
            nhead=n_heads,
            dim_feedforward=c_feedForward,
            dropout=dropout)

        self.transformer = nn.TransformerEncoder(
            self.transformer_layer,
            num_layers=c_layers)

        self.fc_pack_seq_len = nn.Linear(sz_in, 1)
        self.gelu = nn.GELU()
        self.fc_pack_emb = nn.Linear(emb_sz, sz_out)

    def forward(self, x):
        assert x.size(1) == self.sz_in
        x = self.embedding(x) * math.sqrt(self.d_model)
        x = self.pos_encoder(x)
        out = self.transformer(x)
        out = self.fc_pack_seq_len(out.transpose(1,2))[:,:,0]
        out = self.gelu(out)
        out = self.fc_pack_emb(out)
        if self.sz_out > 1:
            return out
        return out[:,0]