add workaround to support fsdp

BiomedSciAI / fuse-med-ml

A python framework accelerating ML based discovery in the medical field by encouraging code reuse. Batteries included :)

Apache License 2.0

134 stars 34 forks source link

add workaround to support fsdp #353

Closed mosheraboh closed 4 months ago

mosheraboh commented 4 months ago

When setting lightning strategy to "fsdp" pytorch_lightning.strategies.FSDPStrategy. model.forward() converts the returned NDict to a simple dictionary, which causes to few issues. This PR, convert it back to NDict if the returned type is not as expected.

mosheraboh commented 4 months ago

Could you please add to the PR's discussion which issue demanded these changes? It's unclear to me and will be useful for future references :)

LGTM

Thanks for the review. Done, hope it's clear now.