Solved by ptrblck in post #2 
              
              
                Yes, it’s possible, but Autograd won’t be able to track these operations, so you would need to write the backward pass manually via a custom autograd.Function as described here.