Hi! I was reading through and familiarizing myself with the codebase.

I noticed that in most files under /datasets, we import Transform like this:
from lobster.transforms import Transform

but in a few cases, the import comes from beignet instead:
from beignet.transforms import Transform

Specifically:

• lobster/src/lobster/datasets/_moleculeace_dataset.py

  Line 7 in c8e9acc

  from beignet.transforms import Transform

• lobster/src/lobster/datasets/_m3_20m_dataset.py

  Line 7 in c8e9acc

  from beignet.transforms import Transform

• lobster/src/lobster/callbacks/_linear_probe_callback.py

  Line 8 in c8e9acc

  from beignet.transforms import Transform

Comparing the two Transform classes in lobster and beignet,
The only implementation difference I found is that the Beignet version explicitly calls zip(..., strict=False), while the Lobster version just calls zip(...) which should default to strict=False.

Proposal:
Would it make sense to standardize the imports by switching these few cases to use lobster.transforms.Transform for consistency?

I’m new to contributing to open source and thought this might be a safe PR to get familiar with the contribution workflow. If this change is welcome, I’d love to submit a PR!