Yes, you can just check the implementations of the optimization algorithms of pytorch (e.g. SGD 2.0k or Adam 943) and use them as a model to help you get started.
Here is an example of an optimization algorithm that was presented at ICLR: AccSGD 1.4k