NameError: name 'str2optimizer8bit_blockwise' is not defined

Hi, thanks for open-sourcing the project. I run the train_network.py module with the default settings and I get the following error.

Traceback (most recent call last):
File "train_network.py", line 773, in
train(args)
File "train_network.py", line 629, in train
optimizer.step()
File "/home/ubuntu/rzamarefat/venv/lib/python3.8/site-packages/accelerate/optimizer.py", line 134, in step
self.scaler.step(self.optimizer, closure)
File "/home/ubuntu/rzamarefat/venv/lib/python3.8/site-packages/torch/cuda/amp/grad_scaler.py", line 370, in step
retval = self._maybe_opt_step(optimizer, optimizer_state, *args, **kwargs)
File "/home/ubuntu/rzamarefat/venv/lib/python3.8/site-packages/torch/cuda/amp/grad_scaler.py", line 290, in _maybe_opt_step
retval = optimizer.step(*args, **kwargs)
File "/home/ubuntu/rzamarefat/venv/lib/python3.8/site-packages/torch/optim/lr_scheduler.py", line 69, in wrapper
return wrapped(*args, **kwargs)
File "/home/ubuntu/rzamarefat/venv/lib/python3.8/site-packages/torch/optim/optimizer.py", line 280, in wrapper
out = func(*args, **kwargs)
File "/home/ubuntu/rzamarefat/venv/lib/python3.8/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
return func(*args, **kwargs)
File "/home/ubuntu/rzamarefat/venv/lib/python3.8/site-packages/bitsandbytes/optim/optimizer.py", line 265, in step
self.update_step(group, p, gindex, pindex)
File "/home/ubuntu/rzamarefat/venv/lib/python3.8/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
return func(*args, **kwargs)
File "/home/ubuntu/rzamarefat/venv/lib/python3.8/site-packages/bitsandbytes/optim/optimizer.py", line 506, in update_step
F.optimizer_update_8bit_blockwise(
File "/home/ubuntu/rzamarefat/venv/lib/python3.8/site-packages/bitsandbytes/functional.py", line 858, in optimizer_update_8bit_blockwise
str2optimizer8bit_blockwise[optimizer_name][0](
NameError: name 'str2optimizer8bit_blockwise' is not defined

Any idea to solve it?

Reproduce the default run of train_network.py and follow the failure from train_network.py lines 629 and 773 into bitsandbytes/functional.py line 858. Check why str2optimizer8bit_blockwise is unavailable during optimizer.step(); done means the default training path completes without this NameError.

Written by the indexing model from the issue text.