Allow missing keys in EmbeddingBagCollection forward #3305
Description
Feature request
Currently, the EmbeddingBagCollection forward pass in torchrec requires all keys to be present in the input. This can be restrictive when some keys are optional or may be missing in certain use cases.
Proposed improvement
Update the EmbeddingBagCollection.forward method to allow missing keys in the input dictionary. If a key is not present, its embedding output could be skipped or returned as a default (e.g., zeros). This will improve flexibility and usability for models with optional features.
Motivation
- Simplifies use cases where not all features are present for every input.
- Reduces the need for preprocessing to add dummy keys.
- Supports more dynamic input scenarios.
Potential implementation
- Modify the forward logic to check for key presence before attempting lookup.
- Document the new behavior and any default values used.
Related issues and discussions
If there are related issues or feature requests, please link them here.
If further details or examples are needed, please specify.