Hi Artix !
You can write your own update function for sure !
To update your weights, you might use the optimiser library. But you
can also do it yourself. For example, you can basically code the
gradient descent, the SGD or Adam using the following code.

net = NN()
learning_rate = 0.01
for param in net.parameters():
weight_update = smth_with_good_dimensions
param.data.sub_(weight_update * learning_rate)

As you can see, you have access to your parameters in net.parameters(), so you can update them like you want.
If you want more specific examples, you can go here where I implemented both SVRG and SAGA : https://github.com/kilianFatras/variance_reduced_neural_networks 664 (variance reduced algorithmes) ! If you have any further question, do not hesitate