Activation Extraction Contains NaN Values After Multiple Runs

[xuemin-yu]

Hello,

I tried to use NeuroX to extract the activations. I recently found that when using neurox.data.extraction.transformers_extractor to extract activations multiple times, subsequent extractions frequently result in activations containing NaN values.

This issue can be observed in the extracted output, as shown in the example below:

{
  "linex_index": 100,
  "features": [
    {
      "token": "<s>",
      "layers": [
        {
          "index": 0,
          "values": [NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, ..., NaN]
        }
      ]
    }
  ]
}