liuyu/qlib
1
0
Fork 0
mirror of https://github.com/microsoft/qlib.git synced 2026-07-06 20:41:09 +08:00
Code Activity

fix TRA when use single head

Browse Source
This commit is contained in:
Dong Zhou
2021-07-22 19:10:17 +08:00
committed by you-n-g
parent 319a2f38cc
commit 4f9f978909
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
qlib/contrib/model/pytorch_tra.py
View File

@@ -526,7 +526,7 @@ class RNN(nn.Module):
if self.use_attn:
laten = self.W(rnn_out).tanh()
scores = self.u(laten).softmax(dim=1)
att_out = (rnn_out * scores).sum(dim=1).squeeze()
att_out = (rnn_out * scores).sum(dim=1)
last_out = torch.cat([last_out, att_out], dim=1)
return last_out
@@ -648,7 +648,7 @@ class TRA(nn.Module):
preds = self.predictors(hidden)
if self.num_states == 1: # no need for router when having only one prediction
return preds.squeeze(-1), preds, None
return preds, None, None
if "TPE" in self.src_info:
out = self.router(hist_loss)[0][:, -1] # TPE