Currently the bottleneck of my model comes from these lines:
outputs = inputs.new_zeros(inputs.size(0), self.in_channels, self.padded_height_in, self.padded_width_in)
outputs[:, :, self.ind_ty, self.ind_tx] = inputs[:, :, self.ind_sy, self.ind_sx]
outputs[:, :, self.ind_up_ty, self.ind_up_tx] = inputs[:, self.ind_up_sc, self.ind_up_sy, self.ind_up_sx]
outputs[:, :, self.ind_down_ty, self.ind_down_tx] = inputs[:, self.ind_down_sc, self.ind_down_sy, self.ind_down_sx]
Is there any faster way of assigning indexed elements from inputs to arbitrary indexed elements of outputs?
I know gather is sometimes faster (specially in the backward pass) than advanced indexing, but I cannot use gather in this case since I am assigning to arbitrary position of outputs instead of to the full outputs.
The pytorch operators taking longer are index and index_put_