Access feature in custom compute_loss method

Iā€™m running into a Hugging Face issue when fine-tuning a model and want to access other features during my custom compute_loss function. However, I am unable to access these in the inputs object passed into the compute_loss function, even though the train_dataset I pass into my CustomTrainer includes this feature. When I call trainer.train(), I get the following error.

---> z = inputs.pop("feature_a")
KeyError: 'feature_a'

I made sure that the train_dataset passed into my CustomTrainer includes ā€œfeature_aā€, as well as include remove_unused_columns=False in the TrainingArguments and implement data_collator that returns {"input_ids": input_ids, "attention_mask": attention_mask, "labels": labels, "feature_a": feature_a}.

Is it possible to access additional features (other than labels and input_ids) in a custom compute_loss function and if so, what I should do to fix this issue of the feature not being in the inputs object passed into compute_loss?