The problem with saving TrainerState (tensors [grad_norm] are inside)
Author: DmitriyLT03
Pull request: no

### System Info


![logs](https://github.com/huggingface/transformers/assets/77717977/3f05ac7d-365e-4a96-a491-8bd6f32b988f)
```
Traceback (most recent call last):
  File "/home/jovyan/scripts/Mistral_7B_trainer_hf_.py", line 106, in <module>
    trainer.train()
  File "/home/user/conda/lib/python3.11/site-packages/transformers/trainer.py", line 1624, in train
    return inner_training_loop(
           ^^^^^^^^^^^^^^^^^^^^
  File "/home/user/conda/lib/python3.11/site-packages/transformers/trainer.py", line 2029, in _inner_training_loop
    self._maybe_log_save_evaluate(tr_loss, grad_norm, model, trial, epoch, ignore_keys_for_eval)
  File "/home/user/conda/lib/python3.11/site-packages/transformers/trainer.py", line 2423, in _maybe_log_save_evaluate
    self._save_checkpoint(model, trial, metrics=metrics)
  File "/home/user/conda/lib/python3.11/site-packages/....../trainer/trainer_hf.py", line 97, in _save_checkpoint
    self.state.save_to_json(os.path.join(ckpt_path, TRAINER_STATE_NAME))
  File "/home/user/conda/lib/python3.11/site-packages/transformers/trainer_callback.py", line 113, in save_to_json
    json_string = json.dumps(dataclasses.asdict(self), indent=2, sort_keys=True) + "\n"
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/conda/lib/python3.11/json/__init__.py", line 238, in dumps
    **kw).encode(obj)
          ^^^^^^^^^^^
  File "/home/user/conda/lib/python3.11/json/encoder.py", line 202, in encode
    chunks = list(chunks)
             ^^^^^^^^^^^^
  File "/home/user/conda/lib/python3.11/json/encoder.py", line 432, in _iterencode
    yield from _iterencode_dict(o, _current_indent_level)
  File "/home/user/conda/lib/python3.11/json/encoder.py", line 406, in _iterencode_dict
    yield from chunks
  File "/home/user/conda/lib/python3.11/json/encoder.py", line 326, in _iterencode_list
    yield from chunks
  File "/home/user/conda/lib/python3.11/json/encoder.py", line 406, in _iterencode_dict
    yield from chunks
  File "/home/user/conda/lib/python3.11/json/encoder.py", line 439, in _iterencode
    o = _default(o)
        ^^^^^^^^^^^
  File "/home/user/conda/lib/python3.11/json/encoder.py", line 180, in default
    raise TypeError(f'Object of type {o.__class__.__name__} '
TypeError: Object of type Tensor is not JSON serializable
```

### Who can help?

@pacman100

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

self.state.save_to_json(os.path.join(ckpt_path, TRAINER_STATE_NAME))

the state has not changed in any way

### Expected behavior

there will be no tensors in log_history and everything will be well preserved in json

Comments:
Hey, can you sahre a reproducible snippet, as well as the output of `transformers-cli env`

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

I am having the same problem.  

Error:
```
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[117], line 1
----> 1 trainer.train()
      2 wandb.finish()

File ~/Repos/advising_gpt/venv/lib/python3.12/site-packages/transformers/trainer.py:1938, in Trainer.train(self, resume_from_checkpoint, trial, ignore_keys_for_eval, **kwargs)
   1936         hf_hub_utils.enable_progress_bars()
   1937 else:
-> 1938     return inner_training_loop(
   1939         args=args,
   1940         resume_from_checkpoint=resume_from_checkpoint,
   1941         trial=trial,
   1942         ignore_keys_for_eval=ignore_keys_for_eval,
   1943     )

File ~/Repos/advising_gpt/venv/lib/python3.12/site-packages/transformers/trainer.py:2356, in Trainer._inner_training_loop(self, batch_size, args, resume_from_checkpoint, trial, ignore_keys_for_eval)
   2353     self.state.epoch = epoch + (step + 1 + steps_skipped) / steps_in_epoch
   2354     self.control = self.callback_handler.on_step_end(args, self.state, self.control)
-> 2356     self._maybe_log_save_evaluate(tr_loss, grad_norm, model, trial, epoch, ignore_keys_for_eval)
   2357 else:
   2358     self.control = self.callback_handler.on_substep_end(args, self.state, self.control)

File ~/Repos/advising_gpt/venv/lib/python3.12/site-packages/transformers/trainer.py:2807, in Trainer._maybe_log_save_evaluate(self, tr_loss, grad_norm, model, trial, epoch, ignore_keys_for_eval)
   2804     metrics = self._evaluate(trial, ignore_keys_for_eval)
   2806 if self.control.should_save:
-> 2807     self._save_checkpoint(model, trial, metrics=metrics)
   2808     self.control = self.callback_handler.on_save(self.args, self.state, self.control)

File ~/Repos/advising_gpt/venv/lib/python3.12/site-packages/transformers/trainer.py:2920, in Trainer._save_checkpoint(self, model, trial, metrics)
   2917 if self.args.should_save:
   2918     # Update the `TrainerControl` state to where we are currently
   2919     self.state.stateful_callbacks["TrainerControl"] = self.control.state()
-> 2920     self.state.save_to_json(os.path.join(output_dir, TRAINER_STATE_NAME))
   2922 if self.args.push_to_hub:
   2923     self._push_from_checkpoint(output_dir)

File ~/Repos/advising_gpt/venv/lib/python3.12/site-packages/transformers/trainer_callback.py:141, in TrainerState.save_to_json(self, json_path)
    139 def save_to_json(self, json_path: str):
    140     """Save the content of this instance in JSON format inside `json_path`."""
--> 141     json_string = json.dumps(dataclasses.asdict(self), indent=2, sort_keys=True) + "\n"
    142     with open(json_path, "w", encoding="utf-8") as f:
    143         f.write(json_string)

File /opt/homebrew/Cellar/python@3.12/3.12.8/Frameworks/Python.framework/Versions/3.12/lib/python3.12/json/__init__.py:238, in dumps(obj, skipkeys, ensure_ascii, check_circular, allow_nan, cls, indent, separators, default, sort_keys, **kw)
    232 if cls is None:
    233     cls = JSONEncoder
    234 return cls(
    235     skipkeys=skipkeys, ensure_ascii=ensure_ascii,
    236     check_circular=check_circular, allow_nan=allow_nan, indent=indent,
    237     separators=separators, default=default, sort_keys=sort_keys,
--> 238     **kw).encode(obj)

File /opt/homebrew/Cellar/python@3.12/3.12.8/Frameworks/Python.framework/Versions/3.12/lib/python3.12/json/encoder.py:202, in JSONEncoder.encode(self, o)
    200 chunks = self.iterencode(o, _one_shot=True)
    201 if not isinstance(chunks, (list, tuple)):
--> 202     chunks = list(chunks)
    203 return ''.join(chunks)

File /opt/homebrew/Cellar/python@3.12/3.12.8/Frameworks/Python.framework/Versions/3.12/lib/python3.12/json/encoder.py:432, in _make_iterencode.<locals>._iterencode(o, _current_indent_level)
    430     yield from _iterencode_list(o, _current_indent_level)
    431 elif isinstance(o, dict):
--> 432     yield from _iterencode_dict(o, _current_indent_level)
    433 else:
    434     if markers is not None:

File /opt/homebrew/Cellar/python@3.12/3.12.8/Frameworks/Python.framework/Versions/3.12/lib/python3.12/json/encoder.py:406, in _make_iterencode.<locals>._iterencode_dict(dct, _current_indent_level)
    404         else:
    405             chunks = _iterencode(value, _current_indent_level)
--> 406         yield from chunks
    407 if newline_indent is not None:
    408     _current_indent_level -= 1

File /opt/homebrew/Cellar/python@3.12/3.12.8/Frameworks/Python.framework/Versions/3.12/lib/python3.12/json/encoder.py:326, in _make_iterencode.<locals>._iterencode_list(lst, _current_indent_level)
    324         else:
    325             chunks = _iterencode(value, _current_indent_level)
--> 326         yield from chunks
    327 if newline_indent is not None:
    328     _current_indent_level -= 1

File /opt/homebrew/Cellar/python@3.12/3.12.8/Frameworks/Python.framework/Versions/3.12/lib/python3.12/json/encoder.py:406, in _make_iterencode.<locals>._iterencode_dict(dct, _current_indent_level)
    404         else:
    405             chunks = _iterencode(value, _current_indent_level)
--> 406         yield from chunks
    407 if newline_indent is not None:
    408     _current_indent_level -= 1

File /opt/homebrew/Cellar/python@3.12/3.12.8/Frameworks/Python.framework/Versions/3.12/lib/python3.12/json/encoder.py:439, in _make_iterencode.<locals>._iterencode(o, _current_indent_level)
    437         raise ValueError("Circular reference detected")
    438     markers[markerid] = o
--> 439 o = _default(o)
    440 yield from _iterencode(o, _current_indent_level)
    441 if markers is not None:

File /opt/homebrew/Cellar/python@3.12/3.12.8/Frameworks/Python.framework/Versions/3.12/lib/python3.12/json/encoder.py:180, in JSONEncoder.default(self, o)
    161 def default(self, o):
    162     """Implement this method in a subclass such that it returns
    163     a serializable object for ``o``, or calls the base implementation
    164     (to raise a ``TypeError``).
   (...)
    178 
    179     """
--> 180     raise TypeError(f'Object of type {o.__class__.__name__} '
    181                     f'is not JSON serializable')

TypeError: Object of type float32 is not JSON serializable
```

`transformers-cli env` output:
```
- `transformers` version: 4.44.2
- Platform: macOS-15.3-arm64-arm-64bit
- Python version: 3.12.8
- Huggingface_hub version: 0.24.0
- Safetensors version: 0.4.3
- Accelerate version: 0.34.2
- Accelerate config: 	not found
- PyTorch version (GPU?): 2.3.1 (False)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: No
```

Code:
```
train_dataset = Dataset.from_dict({
    "sentence": X_train,
    "label": labels
})
model = SentenceTransformer(
    str(MODEL_PATH / "bge-small-finetuned" / "checkpoint-15"),
    device=device,
    model_card_data=SentenceTransformerModelCardData(
        language="en",
        license="mit",
        model_name="bge-small-v1.5 trained on courses with labels"
    ))
loss = BatchHardSoftMarginTripletLoss(model)
args = SentenceTransformerTrainingArguments(
    output_dir = MODEL_PATH / "bge-small-finetuned",
    report_to="wandb",
    logging_steps=5,
    num_train_epochs=100,
    per_device_train_batch_size=8,
    per_device_eval_batch_size=8,
    learning_rate=5e-5,
    warmup_ratio=0,
    fp16=False,
    bf16=True,
    batch_sampler=BatchSamplers.BATCH_SAMPLER,
    eval_strategy="steps",
    eval_steps=25,
    save_strategy="steps",
    save_steps=25,
    use_mps_device=True if device == "mps" else False
)
dev_evaluator = BinaryClassificationEvaluator(
    sentences1=eval_set[0],
    sentences2=eval_set[1],
    labels=eval_set[2],
    name="ppm_toy_eval"
)
trainer = SentenceTransformerTrainer(
    model=model,
    args=args,
    train_dataset=train_dataset,
    loss=loss,
    evaluator=dev_evaluator
)
trainer.train()
```