albanD:

Otherwise, you can try to use something like this  to get a per-sample gradients in one backward pass (more or less) and then do all the post processing after that. That might be faster.


I followed your advise, but to make it work for my GAN application I had to make some major adoptions of the code.  I include this link here for other people’s reference that might encounter the same problem.  Thanks again:

  
      
      GitHub
  
  
    

jusjusjus/autograd-hacks 2

Contribute to jusjusjus/autograd-hacks development by creating an account on GitHub.


  
  
    
    
  
  


Fyi, I also marked it as a solution, even though I think we still have some unresolved questions to the initial observation.