How to get prediction with shape [B, L, N, C]

[xueyagaga]

Great project! However, when I try to get a prediction with shape [B, L, N, C] through STID, I can only get shape [B, L, N, 1]. During the time series embedding process, the input dim C is reduced and unsqueezed with 1. How to keep the output dim the same as C.