ImportError: cannot import name '_refresh_per_optimizer_state' from 'torch.cuda.amp.grad_scaler'

[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

[Guoxt]

"ImportError: cannot import name '_refresh_per_optimizer_state' from 'torch.cuda.amp.grad_scaler'. How can this issue be resolved?"

[mryab]

Hi @Guoxt! This should be fixed in the latest version of Hivemind (by #614), can you upgrade and try again, please?