Closed ZIYU-DEEP closed 1 year ago
Ah, I just realized the cause for the KeyError is that torch.nn.DataParallel(model)
will wrap a module outside the model. A simple fix to avoid this error would be to add one line of model = model.module
before creating layer_collection and calculating FIM.
I wonder if there is any other potential concerns when we use torch.nn.DataParallel(model)
on multiple GPUs with NNGeometry. Thank you!
I am not sure. If you found out, please share the answer :-)
Hey Thomas –
Thank you for creating this terrific package! I am wondering what should be the correct way to compute FIM when we use multiple GPUs with
nn.DataParallel()
to load the network.Specifically, I encountered an KeyError when I tried to run with 3 GPUs and wrap my network with
nn.DataParallel()
. Below is a simplified sample for my code:And I got the following error message:
The error seems not to be raised when using only one GPU.
Would you have any idea how to efficiently solve the issue and compute FIM with multiple GPUs? Thank you so much! : )