Grid Sample Along another axis

Lets say I have a tensor of size [B * 64 * 100 * 100]. Generally, if i wanted to apply a flow field on this tensor along the dimensions 2 and 3, I can use grid sample directly.

But in my case, I have flow field that needs to is shift values in the Channel/64 space. To use my current code directly, could I create a flow image with only change in one direction, rearrange the tensor into [B * 100 * 100 * 64] then apply the field?

1 Like