The shape of torch.tril(eta_mini_batch)

[ZhaoqiZachYang]

In line 963 of the ttt.py, the shape of eta_mini_batch is [B, nh, K, 1], and in line 981, the shape of torch.tril(eta_mini_batch) is ought to be [B,nh,K,K]. However, the output shape of torch.tril([B, nh, K, 1]) is [B, nh, K, 1] rather than [B,nh,K,K]. May I know why the issue happens? Thanks!