I know paradigm and your simple example which is explained in the first hyperlink of my comment. I want to create custom optimizer like this (Custom Optimizer in PyTorch) so I should update weights like this:
weight_update = smth_with_good_dimensions
param.data.sub_(weight_update * learning_rate)
Now, how can we have a function in PyTorch like theano.function for my rms_prop update?