Masking several value inside a tensor with keeping the gradient

Hi, I wonder how to change the value inside a tensor with keeping its gradient.

a = torch.Size([5, 1280, 2])
index = torch.Tensor([100, 200, 300, 400, 500])
r = 5

a[0] = [[x1, y1], [x2, y2], ..., [x_95, y_95], [x_96, y_96], ..., [x_105, y_105], ..., [x1279, y1279]]
a[1] = [[x1, y1], [x2, y2], ..., [x_195, y_195], [x_196, y_196], ..., [x_205, y_205], ..., [x1279, y1279]]
a[2] = [[x1, y1], [x2, y2], ..., [x_295, y_295], [x_296, y_296], ..., [x_305, y_305], ..., [x1279, y1279]]
a[3] = [[x1, y1], [x2, y2], ..., [x_395, y_395], [x_396, y_396], ..., [x_405, y_405], ..., [x1279, y1279]]
a[4] = [[x1, y1], [x2, y2], ..., [x_495, y_495], [x_496, y_496], ..., [x_505, y_505], ..., [x1279, y1279]]

It means that every index[i] indicates the index of 5 a[i].
The resulting tensor is expected like below

a[0] = [[0, 0], [0, 0], ..., [x_95, y_95], [x_96, y_96], ..., [x_105, y_105], ..., [0, 0]]
a[1] = [[0, 0], [0, 0], ..., [x_195, y_195], [x_196, y_196], ..., [x_205, y_205], ..., [0, 0]]
a[2] = [[0, 0], [0, 0], ..., [x_295, y_295], [x_296, y_296], ..., [x_305, y_305], ..., [0, 0]]
a[3] = [[0, 0], [0, 0], ..., [x_395, y_395], [x_396, y_396], ..., [x_405, y_405], ..., [0, 0]]
a[4] = [[0, 0], [0, 0], ..., [x_495, y_495], [x_496, y_496], ..., [x_505, y_505], ..., [0, 0]]

Any help? Please comment on any unclear explanations.

Thanks,