Судя по тому, как это выглядит на приведенном вами рисунке, они добавляют априорные значения местоположения к данным, т.е.
location_priors = generate_gaussians(positions, variances, data.size())
data_w_loc_priors = T.cat((data, location_priors), dim=1)
Теперь количество in_channels для вашей свертки необходимо соответствующим образом скорректировать: раньше у вас было 512 in_channels, теперь у вас 512 + количество априорных мест.