Warning bieng raised with 'cpu-offloading' option and generating structure takes forever

[amani-acog]

NOTE: if this is not a bug report, please use the GitHub Discussions for support questions (How do I do X?), feature requests, ideas, showcasing new applications, etc.

Bug description
I am running esmfold in a docker container with GPUs enabled (not fully available gpu memory). I am getting the following warning:
UserWarning: Module is put on CPU and will thus have flattening and sharding run on CPU, which is less efficient than on GPU. We recommend passing in device_id argument which will enable FSDP to put module on GPU device, module must also be on GPU device to work with sync_module_states=True flag which requires GPU communication.
"Module is put on CPU and will thus have flattening and sharding"

However the there is no error and execution seems to be going on, but there is no trace of CPU or GPU usage. It did not generate predictions either.

Reproduction steps

Expected behavior
Give a clear and concise description of what you expected to happen.

Logs
Please paste the command line output:

Output goes here

Additional context
Add any other context about the problem here. (like proxy settings, network setup, overall goals, etc.)