Solved by ptrblck in post #2 
              
              
                You could probably directly index the mapping as seen here: 
# create input
x = torch.arange(2*3*32*32).to(torch.uint8)
x = x.view(2, 3, 32, 32)
print(x)

# create mapping by shifting it by 1
mapping = torch.arange(0, 256, dtype=torch.uint8) + 1
# mapping[index] would contain the new value

# index …