Adding comments to test_cuda.py

pytorch · May 14, 2024 · a0af2fe · a0af2fe
1 parent d4815a7
commit a0af2fe
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/test/test_cuda.py b/test/test_cuda.py
@@ -4465,7 +4465,10 @@ def test_graph_optims(self, device, dtype, optim_info):
 
         for optim_input in all_optim_inputs:
             kwargs = optim_input.kwargs
-            kwargs["lr"]=0.1
+
+            # lr as a Tensor is not supported when capturable=False and foreach=True for torch.optim.adam
+            # and torch.optim.adamw
+            kwargs["lr"] = 0.1 
 
             for actually_do_graphs in (True, False):
                 params = [
@@ -4480,7 +4483,6 @@ def test_graph_optims(self, device, dtype, optim_info):
                 ]
 
                 # Control (capturable=False)
-
                 kwargs["capturable"] = False
 
                 opt = optim_cls(params_control, **kwargs)