Assign a value to tensor components based on condition

How can I do this?
I have a tensor A of size, say, 4x10x64x64.
I want to find all elements where A is bigger than 255, and set them to 255.

I did idx = torch.gt(A, 255), and obtained another tensor which contains ones or zeros. But I don’t understand how to pass these to A to alter it.

1 Like

torch.clamp(A, max=255)

1 Like

Thanks! That did the trick)