Hey @albanD, I am trying out SGD and it is working. Thank you for the suggestion. A question though. If I am optimizing over a huge tensor with Adam but am only operating on a narrowed view of the tensor at a time. Will the optimization procedure still update (or allocate memory for updating) all the gradients? Can I change this behavior in any way?