Hi everyone,
I have a specific use-case wherein I want to dynamically move model parameters to GPU from CPU, and once the weight updates are done (on the GPU), transfer it back again on the CPU.
To be specific, I am dealing with a classification problem, where the number of classes is huge. Let’s say we have a linear setup, i.e we want to learn a weight matrix of size [num_classes x feature_size]. The problem is that we can’t afford to keep this matrix completely on the GPU. Hence, what I wanted is to initialize everything on CPU and keep loading parts of the matrix onto GPU and update correspondingly.
With this ideology, I’ll have to declare the weight parameters on CPU (along with their optimizers) and when I want to load part of the matrix on GPU, I’ll have to call .cuda(), which, as the docs suggest (LINK 2) needs to be called before optimizer construction. Anyways, on trying this approach, as expected the model weights don’t update.
Any help would be greatly appreciated.
Thanks,
Noveen Sachdeva