mirror of
https://github.com/microsoft/qlib.git
synced 2026-07-17 17:34:35 +08:00
fix TRA when use single head
This commit is contained in:
@@ -526,7 +526,7 @@ class RNN(nn.Module):
|
|||||||
if self.use_attn:
|
if self.use_attn:
|
||||||
laten = self.W(rnn_out).tanh()
|
laten = self.W(rnn_out).tanh()
|
||||||
scores = self.u(laten).softmax(dim=1)
|
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)
|
last_out = torch.cat([last_out, att_out], dim=1)
|
||||||
|
|
||||||
return last_out
|
return last_out
|
||||||
@@ -648,7 +648,7 @@ class TRA(nn.Module):
|
|||||||
preds = self.predictors(hidden)
|
preds = self.predictors(hidden)
|
||||||
|
|
||||||
if self.num_states == 1: # no need for router when having only one prediction
|
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:
|
if "TPE" in self.src_info:
|
||||||
out = self.router(hist_loss)[0][:, -1] # TPE
|
out = self.router(hist_loss)[0][:, -1] # TPE
|
||||||
|
|||||||
Reference in New Issue
Block a user