With FSDP, when using model's forward in on_fit_end() hook weights are misshappen #20254
Unanswered
QuentinAndre11
asked this question in
code help: CV
Replies: 0 comments
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
-
Hi everyone !
I am training an image classifier and would like to see the embeddings at the end of training, but I don't find how to do it while using FSDP, since the weights seem to get flattenned outside of train/test/validation/_step. Indeed, with the following code, I get a RuntimeError: weight should have at least three dimensions.
I use pytorch v2.1.0 with CUDA 118 and lightning v2.3.3
Thus, how should I do forward calls on_fit_end() hook while using FSDP? Or is there a better way to get the embeddings (and if so, why shouldn't I use the weights outside of training steps)?
Thanks for your help already
Beta Was this translation helpful? Give feedback.
All reactions