Open jackdent opened 2 months ago
The Lightning apply_to_collection
logic is defined here and relies on dataclass.fields
, which doesn't include cached properties
@awaelchli, do you have any experience with this one?
Hey @jackdent
This is a rare use case and I won't have the bandwidth to look into it. We would be grateful for a contribution here if you're interested. The fix is probably to just reset the cache when running apply_to_collection
.
Motivation
When running
apply_to_collection
on adataclass
, cached properties do not get modified. This can cause subtle issues: for example, suppose I initialize adataclass
on CPU in a dataworker, and then move it onto GPU for a model batch. All of thedataclass
fields that containTensor
s get moved correctly, but thecached_property
s continue to residue on the original device.Steps to reproduce
Yields the following output: