Fix `orientation` batching

a-r-j / ProteinWorkshop

Benchmarking framework for protein representation learning. Includes a large number of pre-training and downstream task datasets, models and training/task utilities. (ICLR 2024)

https://proteins.sh/

MIT License

194 stars 16 forks source link

Fix `orientation` batching #58

Closed amorehead closed 9 months ago

amorehead commented 10 months ago

Fixes orientation feature batching for vector-equivariant methods.
Without this fix, for vector-equivariant methods that use such features, precisely batch.num_graphs - 2 nodes in each of their batched protein graphs would have their vector node input features normalized to a nonsense (invalid) vector.

a-r-j commented 9 months ago

@amorehead Happy for me to merge?

amorehead commented 9 months ago

@arian, yep, should be good to go. Let me know if you run into any unexpected edge cases while training. I believe this should handle all of them.