Say I have a tensor A: [B, n, 64, 64]. Now I want to have n GRUCell units process this tensor across dim=1.
That is GRUCell_0 gets A[:, 0] and so on and so forth.
Is there a parallel way to achieve this behavior instead of a for loop?
Thanks!
Say I have a tensor A: [B, n, 64, 64]. Now I want to have n GRUCell units process this tensor across dim=1.
That is GRUCell_0 gets A[:, 0] and so on and so forth.
Is there a parallel way to achieve this behavior instead of a for loop?
Thanks!