Is there an easy way to assign values to a tensor in a way that’s similar to torch.select?

For instance, say I have a tensor with n dimensions and I want to change the mth “row” of the kth dimension to be 1. If n = 5, m = 2, and k = 3, I could do something like tensor[:, :, :, 2, :] = 1 - but let’s say that m, n, and k are not known beforehand, or I want to loop through different values of m and k, etc. If I just wanted to retrieve the values of these slices, torch.select(tensor, 3, 2) would work, but I want to be able to make changes to this slice.

I’m hoping there’s a way to somehow loop through the dimensions, sort of like

for dim in range(len(tensor.shape)):
torch.select(tensor, dim, 0) *= 0.5 # But with a different function than torch.select
torch.select(tensor, dim, -1) *= 0.5