Index Error Invalid Ket is out of bounds for size 0 for code-search-net/code_search_net dataset

[harshakhmk]

Describe the bug

I am trying to do model pruning on sentence-transformers/all-mini-L6-v2 for the code-search-net/code_search_net dataset using INCTrainer class
However I am getting below error

raise IndexError(f"Invalid Key: {key is our of bounds for size {size}")
IndexError: Invalid key: 1840208 is out of bounds for size 0

Steps to reproduce the bug

Model pruning on the above dataset using the below guide

https://huggingface.co/docs/optimum/en/intel/neural_compressor/optimization#pruning

Expected behavior

The modsl should be successfully pruned

Environment info

Torch version: 2.4.1
Python version: 3.8.10