How to set values at tensor at dynamic dim?

1

How to set values at tensor at dynamic dim? The sizes of the tensor parameter can be dynamic.
For example, for tensor = torch.rand(size=(4, 5, 6)), how to create a function like below? :

def ff(tensor, dim, pos, value):
    # this function can also be applied to tensors 
    # with shapes like [7,5,6,8], [11, 5, 7,5,6,8], etc.
    if dim == 0:
        tensor[ pos, :, :] = value
    elif dim == 1:
        tensor[:, pos, :] = value
    elif dim == 2:
        tensor[:, :, pos] = value
2

Depending on the operation, you could e.g. use tensor.index_copy_, which accepts a dim argument.