Как правильно запустить претренированную модель с pytorch?

Question

[Ася]

Хочу использовать претренированную модель отсюда

Но вылетает такая ошибка:

Missing key(s) in state_dict: "gru_1.weight_ih_l1", "gru_1.weight_hh_l1", "gru_1.bias_ih_l1", "gru_1.bias_hh_l1", "gru_1.weight_ih_l2", "gru_1.weight_hh_l2", "gru_1.bias_ih_l2", "gru_1.bias_hh_l2". 
Unexpected key(s) in state_dict: "gru_2.weight_ih_l0", "gru_2.weight_hh_l0", "gru_2.bias_ih_l0", "gru_2.bias_hh_l0", "gru_3.weight_ih_l0", "gru_3.weight_hh_l0", "gru_3.bias_ih_l0", "gru_3.bias_hh_l0".

когда пытаюсь

decoder = MolDecoder(c=len(charset))
    # decoder = WrappedModel(decoder)
    decoder.apply(initialize_weights)
    checkpoint = torch.load('./ref/ref_model.pth.tar',  map_location={"cuda:0" : "cpu"})
    decoder.load_state_dict(checkpoint['decoder'])

class MolDecoder(nn.Module):

    def __init__(self, i=292, o=120, c=35):
        super(MolDecoder, self).__init__()

        self.latent_input = nn.Sequential(nn.Linear(i, i),
                                          SELU(inplace=True))
        self.repeat_vector = Repeat(o)
        self.gru_1 = nn.GRU(i, 501, 3, batch_first=True)
        # self.gru_2 = nn.GRU(501, 501, 3, batch_first=True)
        # self.gru_3 = nn.GRU(501, 501, 3, batch_first=True)
        self.decoded_mean = TimeDistributed(nn.Sequential(nn.Linear(501, 55), nn.Softmax()))

    def forward(self, x):
        out = self.latent_input(x)
        out = self.repeat_vector(out)
        out, h = self.gru_1(out)
        return self.decoded_mean(out)

Как исправить? запускаю на cpu