Fill tensor with another tensor where mask is true

I need to insert elements of tensor new into a tensor old with a certain probability, let’s say that it is 0.8 for simplicity. Substantially this is what masked_fill would do, but it only works with monodimensional tensor. Actually I am doing

    prob = torch.rand(trgs.shape, dtype=torch.float32).to(trgs.device)
    mask = prob < 0.8

    dim1, dim2, dim3, dim4 = new.shape
    for a in range(dim1):
        for b in range(dim2):
            for c in range(dim3):
                for d in range(dim4):
                    old[a][b][c][d] = old[a][b][c][d] if mask[a][b][c][d] else new[a][b][c][d]

which is awful. I would like something like

    prob = torch.rand(trgs.shape, dtype=torch.float32).to(trgs.device)
    mask = prob < 0.8

    old = trgs.multidimensional_masked_fill(mask, new)
def old_imp(old, new, mask):
    dim1, dim2, dim3, dim4 = new.shape
    for a in range(dim1):
        for b in range(dim2):
            for c in range(dim3):
                for d in range(dim4):
                    old[a][b][c][d] = old[a][b][c][d] if mask[a][b][c][d] else new[a][b][c][d]
    return old 

def new_imp(old, new, mask):
    new[mask] = old[mask]
    return new

old = torch.rand(2, 3, 4, 5)
new = torch.rand(2, 3, 4, 5)
prob = torch.rand(old.shape, dtype=torch.float32)
mask = prob < 0.8

old_res = old_imp(old.clone(), new.clone(), mask.clone())
new_res = new_imp(old.clone(), new.clone(), mask.clone())
print(old_res ==  new_res)