upon trying your code, it seems that sorting before unique is eating some time. unique takes a sorted argument of its own that takes less time (i think sorting before unique runs 2 passes through the data, while sorting during unique does it in 1 pass). also you can use the inverse indexes returned by unique to simplify the code:
import torch
import time
W, H = 3, 3
image_flatten = torch.tensor([2, 3, 4, 1, 6, 2, 2, 5, 3])
print(image_flatten.view(W, H))
# sorted_image_flatten, idx_sort = torch.sort(image_flatten)
unique_values, inverse_indexes, count = torch.unique(image_flatten, sorted = True,return_counts=True, return_inverse=True)
non_repeated_vals = (count == 1).float()
mask_flatten = non_repeated_vals[inverse_indexes]
mask = mask_flatten.view(H,W)