The additional dimension in weight_flopk_prev might cause some trouble.
Here is a small example:
weight = torch.randn(10, 6, 3, 3)
idx0 = torch.randint(0, 10, (5,))
idx1 = torch.randint(0, 6, (3,))

new_weight = weight[idx0.unsqueeze(1), idx1]
print(new_weight.shape)
> torch.Size([5, 3, 3, 3])

Also, if you are trying to manipulate the parameters directly, warp it in a with torch.no_grad() block and don’t use .data as this might yield to wrong calculations.