🐛 Describe the bug

When the original LR is a tensor, _init_optim_state() should respect the tensor-ness of LR. This probably doesn't matter for built-in PyTorch optimizers, but for torchao's low-bit optimizers, LR is expected to be a tensor. There will be an error otherwise. Related to pytorch/ao#1189

I propose to change L614 to

param_group["lr"] = torch.tensor(0.0) if isinstance(param_group["lr"], torch.Tensor) else 0.0

cc: @awgu

Versions

2.6.0.dev20241029

cc @H-Huang @awgu @kwen2501 @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @c-p-i-o @ezyang @chauhang @penguinwu