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)