Repository: pytorch
Issue: Incomplete check of LR as a tensor in Optimizer
Labels: module: optimizer,triaged,actionable

### 🐛 Describe the bug

A tutorial "[Running the compiled optimizer with an LR Scheduler](https://pytorch.org/tutorials/recipes/compiling_optimizer_lr_scheduler.html)" presents how LR as a tensor is used. According to this tutorial, we should use a 0-dim tensor for LR. However, Optimizer can accept a 1-dim tensor of size 1. When using a 1-dim tensor, we get a runtime error.

```
optimizer = torch.optim.SGD(model.parameters(), lr=torch.tensor([0.01]).cuda())
scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, total_iters=5)
```

```
RuntimeError: fill_ only supports 0-dimension value tensor but got tensor with 1 dimensions.
```

Currently, Optimizer checks the tensor size:

https://github.com/pytorch/pytorch/blob/faa10faa2cad1cf6eef95a3e5bda255be6bc4c87/torch/optim/adam.py#L50-L56

I think that Optimizer should check the tensor dimension:

```
     if lr.dim() != 0: 
         raise ValueError("Tensor lr must be 0-dimension") 
```

The whole code I used:
```
import torch
print(torch.__version__)

# exit cleanly if we are on a device that doesn't support ``torch.compile``
if torch.cuda.get_device_capability() < (7, 0):
    print("Exiting because torch.compile is not supported on this device.")
    import sys
    sys.exit(0)

# Setup random seed
torch.manual_seed(12345)

# Create simple model
model = torch.nn.Sequential(
    *[torch.nn.Linear(1024, 1024, False, device="cuda") for _ in range(10)]
)
input = torch.rand(1024, device="cuda")

# run forward pass
output = model(input)

# run backward to populate the grads for our optimizer below
output.sum().backward()

optimizer = torch.optim.SGD(model.parameters(), lr=torch.tensor([0.01]).cuda())
scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, total_iters=5)

def fn():
    optimizer.step()
    scheduler.step()

opt_fn = torch.compile(fn)

for i in range(5):    
    opt_fn()
    print(i, optimizer.param_groups[0]['lr'].item())
```


### Error logs

```
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-2-67dacbf46b2f> in <cell line: 26>()
     24 
     25 optimizer = torch.optim.SGD(model.parameters(), lr=torch.tensor([0.01]).cuda())
---> 26 scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, total_iters=5)
     27 
     28 def fn():

/usr/local/lib/python3.10/dist-packages/torch/optim/lr_scheduler.py in __init__(self, optimizer, start_factor, end_factor, total_iters, last_epoch, verbose)
    746         self.end_factor = end_factor
    747         self.total_iters = total_iters
--> 748         super().__init__(optimizer, last_epoch, verbose)
    749 
    750     def get_lr(self):

/usr/local/lib/python3.10/dist-packages/torch/optim/lr_scheduler.py in __init__(self, optimizer, last_epoch, verbose)
    144         patch_track_step_called(self.optimizer)
    145         self.verbose = _check_verbose_deprecated_warning(verbose)
--> 146         self._initial_step()
    147 
    148     def _initial_step(self):

/usr/local/lib/python3.10/dist-packages/torch/optim/lr_scheduler.py in _initial_step(self)
    149         """Initialize step counts and perform a step."""
    150         self._step_count = 0
--> 151         self.step()
    152 
    153     def state_dict(self):

/usr/local/lib/python3.10/dist-packages/torch/optim/lr_scheduler.py in step(self, epoch)
    248             param_group, lr = data
    249             if isinstance(param_group["lr"], Tensor):
--> 250                 param_group["lr"].fill_(lr)
    251             else:
    252                 param_group["lr"] = lr

RuntimeError: fill_ only supports 0-dimension value tensor but got tensor with 1 dimensions.
```

### Versions

PyTorch version: 2.5.1+cu121
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.3 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: 14.0.0-1ubuntu1.1
CMake version: version 3.31.2
Libc version: glibc-2.35

Python version: 3.10.12 (main, Nov  6 2024, 20:22:13) [GCC 11.4.0] (64-bit runtime)
Python platform: Linux-6.6.56+-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: 12.2.140
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: Tesla T4
GPU 1: Tesla T4

Nvidia driver version: 560.35.03
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.9.6
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                         x86_64
CPU op-mode(s):                       32-bit, 64-bit
Address sizes:                        46 bits physical, 48 bits virtual
Byte Order:                           Little Endian
CPU(s):                               4
On-line CPU(s) list:                  0-3
Vendor ID:                            GenuineIntel
Model name:                           Intel(R) Xeon(R) CPU @ 2.00GHz
CPU family:                           6
Model:                                85
Thread(s) per core:                   2
Core(s) per socket:                   2
Socket(s):                            1
Stepping:                             3
BogoMIPS:                             4000.30
Flags:                                fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc rep_good nopl xtopology nonstop_tsc cpuid tsc_known_freq pni pclmulqdq ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch pti ssbd ibrs ibpb stibp fsgsbase tsc_adjust bmi1 hle avx2 smep bmi2 erms invpcid rtm mpx avx512f avx512dq rdseed adx smap clflushopt clwb avx512cd avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves arat md_clear arch_capabilities
Hypervisor vendor:                    KVM
Virtualization type:                  full
L1d cache:                            64 KiB (2 instances)
L1i cache:                            64 KiB (2 instances)
L2 cache:                             2 MiB (2 instances)
L3 cache:                             38.5 MiB (1 instance)
NUMA node(s):                         1
NUMA node0 CPU(s):                    0-3
Vulnerability Gather data sampling:   Not affected
Vulnerability Itlb multihit:          Not affected
Vulnerability L1tf:                   Mitigation; PTE Inversion
Vulnerability Mds:                    Mitigation; Clear CPU buffers; SMT Host state unknown
Vulnerability Meltdown:               Mitigation; PTI
Vulnerability Mmio stale data:        Vulnerable: Clear CPU buffers attempted, no microcode; SMT Host state unknown
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed:               Mitigation; IBRS
Vulnerability Spec rstack overflow:   Not affected
Vulnerability Spec store bypass:      Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:             Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:             Mitigation; IBRS; IBPB conditional; STIBP conditional; RSB filling; PBRSB-eIBRS Not affected; BHI SW loop, KVM SW loop
Vulnerability Srbds:                  Not affected
Vulnerability Tsx async abort:        Mitigation; Clear CPU buffers; SMT Host state unknown

Versions of relevant libraries:
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.26.4
[pip3] nvidia-cublas-cu12==12.6.4.1
[pip3] nvidia-cuda-cupti-cu12==12.6.80
[pip3] nvidia-cuda-runtime-cu12==12.6.77
[pip3] nvidia-cudnn-cu12==9.6.0.74
[pip3] nvidia-cufft-cu12==11.3.0.4
[pip3] nvidia-curand-cu12==10.3.7.77
[pip3] nvidia-cusolver-cu12==11.7.1.2
[pip3] nvidia-cusparse-cu12==12.5.4.2
[pip3] nvidia-nccl-cu12==2.23.4
[pip3] nvidia-nvjitlink-cu12==12.6.85
[pip3] nvtx==0.2.10
[pip3] onnx==1.17.0
[pip3] optree==0.13.1
[pip3] pynvjitlink-cu12==0.4.0
[pip3] pytorch-ignite==0.5.1
[pip3] pytorch-lightning==2.5.0.post0
[pip3] torch==2.5.1+cu121
[pip3] torchaudio==2.5.1+cu121
[pip3] torchinfo==1.8.0
[pip3] torchmetrics==1.6.1
[pip3] torchsummary==1.5.1
[pip3] torchtune==0.5.0
[pip3] torchvision==0.20.1+cu121
[conda] Could not collect

cc @vincentqb @jbschlosser @albanD @janeyx99 @crcrpar @chauhang @penguinwu