Don't understand VIN last step

[QiXuanWang]

    slice_s1 = S1.long().expand(config.imsize, 1, config.l_q, q.size(0))
    slice_s1 = slice_s1.permute(3, 2, 1, 0)
    q_out = q.gather(2, slice_s1).squeeze(2)

What does this 3 lines do?

[kentsommer]

Hi @QiXuanWang this is to achieve the following which is described in the author's original implementation: https://github.com/avivt/VIN/blob/master/vin.py#L242

# Select the conv-net channels at the state position (S1,S2).
# This intuitively corresponds to each channel representing an action, and the convnet the Q function.
# The tricky thing is we want to select the same (S1,S2) position *for each* channel and for each sample