Thursday, 29 November 2018

Index pytorch 4d tensor by values in 2d tensor

I have two pytorch tensors:

  • X with shape (A, B, C, D)
  • I with shape (A, B)

Values in I are integers in range [0, C).


What is the most efficient way to get tensor Y with shape (A, B, D), such that:

Y[i][j][k] = X[i][j][ I[i][j] ][k]



from Index pytorch 4d tensor by values in 2d tensor

No comments:

Post a Comment