Hi, I am transfering Style Swap(based on torch) into pytorch. I wonder if this could be improved.
Input is 4 dimensional, 1*C*H*W.
  _, argmax = torch.max(input,1)
      for i=1,self.output:size(3) do
        for j=1,self.output:size(4) do
          ind = argmax[{1,1,i,j}]
          self.output[{1,ind,i,j}] = 1
        end
     end
     return self.output
end


Can this be improved in the lua version? So I can pull a request to the original repo.
How to improved it in pytorch version.
BTW, If I want to record the ind. How to do it efficiently?

--addtional
local spSize = self.output:size(3)*self.output:size(4)
self.ind_vec = torch.Tensor(spSize):zero()

 _, argmax = torch.max(input,1)
 for i=1,self.output:size(3) do
     for j=1,self.output:size(4) do
         ind = argmax[{1,1,i,j}]
         self.output[{1,ind,i,j}] = 1
         -- additional
         tmp = (i-1)*self.output:size(3)+j
         self.ind_vec[tmp] = ind
    end
  end
return self.output
end