move to torch.optim for optimization

[Mayukhdeb]

Currently, the optimization is done using a simple implementation which goes as:

image_tensor.data = image_tensor.data + lr *(gradients_tensor.data /grad_norm)

to something like

optim = optim.some_optimizer(image_tensor.parameters(), lr= config["learning_rate"], momentum=0.9)
loss.backward()

Have a loot at the shampoo optimizer, it can be used for "preconditioning"

Two advantages:

• Cleaner code
• (maybe) better performance