Closed johann-petrak closed 5 years ago
OK, turns out the embeddings parameters are correctly on the CPU, but the batch indices are on the GPU.
Ahhh - the problem is that the flag that caches the CUDA state gets saved with the module.
OK, this can be fixed by overriding __getstate__
to avoid saving the transient variable and/or by overriding __setstate__
to initialise the transient variable on unpickling.
Also, instead of moving everything back to the cpu after loading, we could try to re-map to the CPU (or desired GPU) already during load time using
torch.load('mymodule', map_location=lambda storage, location: 'cpu')
Should implement the following changes:
Still missing the wrapper load parameter functionality
Finally fixed as of a187946f4
When restoring the model, it seems to get moved automatically to CUDA at load time. Even moving the top module to the CPU does not seem to help - eventually the embeddings parameters are still on the GPU.