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

### 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