Sorry, I’ve just posted a solution to another problem and it has been withdrawn.
To accelerate it, I use advance indexing.(I haven’t finished complete perf comparison. But it does run much faster for my case.) I only use batchsize=1 in my training. So I haven’t work out a plan to handle multi batchsize.
import torch
import datetime
import time
import numpy

bs = 1  # has to be 1 for now.
c = 256
h = 32
w = 32

input  = torch.randn(1,c,h,w).mul_(10).floor()


sp_y = torch.arange(0,w).long()
sp_y = torch.cat([sp_y]*h)

lst = []
for i in range(h):
    lst.extend([i]*w)
sp_x = lst
sp_x = torch.from_numpy(numpy.array(sp_x))


print('Method 1')
start1 = time.time()
_,c_max = torch.max(input, 1)

c_max_flatten = c_max.view(-1)

input_zero = torch.zeros_like(input)


input_zero[:,c_max_flatten,sp_x,sp_y]=1
indlst1 = c_max_flatten
end1 = time.time()

print(type(sp_x))
print(type(sp_y))



input_zero=input_zero.cuda()




print('Method 2')

input_zero2 = torch.zeros_like(input).cuda()
indlst2 = torch.zeros(h*w).long()
start2 = time.time()
_, arg_max = torch.max(input,1,keepdim=True)
for i in range(h):
    for j in range(w):
        ind = arg_max[:,0,i,j]
        input_zero2[:,ind[0],i,j] = 1
        tmp = i*w+j
        indlst2[tmp] = ind[0]

end2 = time.time()

print('Speedup ratio:')
print((end2-start2)/(end1-start1)) # about 1-15 times faster
print('Before:')
print(end2-start2)
print('After:')
print(end1-start1)


print('error:')
print(torch.sum(indlst1-indlst2))
print(torch.sum(input_zero-input_zero2))