Replace part of (with random index) a multi-dimensional tensor with another tensor


I have two tensors x and y, both with size (a * b * c). Now I want to do this:

def forward(x, y)
    a, b, c = x.shape
    for i in range(b):
        rand_idx = np.random.randint(a)
        x[rand_idx:, i, :] = y[rand_idx:, i, :]
    return x

and make sure the gradient graph doesn’t break.
Is there a way in PyTorch to write this without the for loop?