ImportError: cannot import name '_refresh_per_optimizer_state' from 'torch.cuda.amp.grad_scaler' (/home/guoxt/miniconda3/lib/python3.12/site-packages/torch/cuda/amp/grad_scaler.py)

[Guoxt]

(base) guoxt@dgx1-11:$ hivemind-dht --identity_path fixed_private_key.pem --host_maddrs /ip4/0.0.0.0/tcp/30001
Traceback (most recent call last):
File "/home/guoxt/miniconda3/bin/hivemind-dht", line 5, in
from hivemind.hivemind_cli.run_dht import main
File "/home/guoxt/miniconda3/lib/python3.12/site-packages/hivemind/init.py", line 12, in
from hivemind.optim import GradScaler, Optimizer, TrainingAverager
File "/home/guoxt/miniconda3/lib/python3.12/site-packages/hivemind/optim/init.py", line 1, in
from hivemind.optim.grad_scaler import GradScaler, HivemindGradScaler
File "/home/guoxt/miniconda3/lib/python3.12/site-packages/hivemind/optim/grad_scaler.py", line 8, in
from torch.cuda.amp.grad_scaler import OptState, _refresh_per_optimizer_state
ImportError: cannot import name '_refresh_per_optimizer_state' from 'torch.cuda.amp.grad_scaler' (/home/guoxt/miniconda3/lib/python3.12/site-packages/torch/cuda/amp/grad_scaler.py)
(base) guoxt@dgx1-11:$ python -c "import torch; print(torch.version)"
2.6.0+cu124
(base) guoxt@dgx1-11:~$ pip install torch==1.10.0