From 5141ae69e9f87ef3fbc62a95c2a19adbacfc1a6e Mon Sep 17 00:00:00 2001 From: "Dr. Kashif Rasul" Date: Tue, 28 Jan 2020 10:50:53 +0100 Subject: [PATCH] fix loss --- pts/model/transformer/transformer_network.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/pts/model/transformer/transformer_network.py b/pts/model/transformer/transformer_network.py index e23d502..9fa02f7 100644 --- a/pts/model/transformer/transformer_network.py +++ b/pts/model/transformer/transformer_network.py @@ -55,9 +55,11 @@ class TransformerNetwork(nn.Module): self.target_shape = distr_output.event_shape + # [B, T, input_size] -> [B, T, d_model] self.encoder_input = nn.Linear(input_size, d_model) self.decoder_input = nn.Linear(input_size, d_model) + # [B, T, d_model] where d_model / num_heads... self.transformer = nn.Transformer( d_model=d_model, nhead=num_heads, @@ -289,7 +291,7 @@ class TransformerTrainingNetwork(TransformerNetwork): # compute loss distr_args = self.proj_dist_args(dec_output) distr = self.distr_output.distribution(distr_args, scale=scale) - loss = distr.loss(future_target) + loss = - distr.log_prob(future_target) return loss.mean()