More assertive require_grad check (#4709)

Fixes pytorch/test-infra#4687

Author: kit1980

tests/fixtures/misc/checker/require_grad.py

@@ -1,9 +1,11 @@
+import random
 import torch
 x = torch.zeros(1)
 x.require_grad = False
 x.require_grad = True
+grad = random.choice([False, True])
+x.require_grad = grad
 
 # Don't trigger
 x.requires_grad = False
 require_grad = False
-x.require_grad = 10

tests/fixtures/misc/checker/require_grad.txt

@@ -1,2 +1,3 @@
-3:1 TOR002 Likely typo `require_grad` in assignment. Did you mean `requires_grad`?
 4:1 TOR002 Likely typo `require_grad` in assignment. Did you mean `requires_grad`?
+5:1 TOR002 Likely typo `require_grad` in assignment. Did you mean `requires_grad`?
+7:1 TOR002 Likely typo `require_grad` in assignment. Did you mean `requires_grad`?

tests/fixtures/misc/codemod/require_grad.py

@@ -2,3 +2,9 @@
 x = torch.zeros(1)
 x.require_grad = False
 x.require_grad = True
+
+# from https://github.com/pytorch/test-infra/issues/4687
+import torch.nn as nn
+model = nn.Module()
+for name, param in model.named_parameters():
+    param.require_grad = 'specific_layer' in name

tests/fixtures/misc/codemod/require_grad.py.out

@@ -2,3 +2,9 @@ import torch
 x = torch.zeros(1)
 x.requires_grad = False
 x.requires_grad = True
+
+# from https://github.com/pytorch/test-infra/issues/4687
+import torch.nn as nn
+model = nn.Module()
+for name, param in model.named_parameters():
+    param.requires_grad = 'specific_layer' in name

torchfix/visitors/misc/__init__.py

@@ -14,12 +14,10 @@ class TorchRequireGradVisitor(TorchVisitor):
     MESSAGE = "Likely typo `require_grad` in assignment. Did you mean `requires_grad`?"
 
     def visit_Assign(self, node):
-        # Look for any assignment with `require_grad` attribute on the left
-        # and `False` or `True` on the right.
+        # Look for any assignment with `require_grad` attribute on the left.
         #
-        # If this causes false-positives on real code (unlikely),
-        # we can do type inference (not sure if feasible here) or
-        # at least check that `torch` is imported in the file.
+        # This is unlikely to cause false-positives on real code, especially
+        # because TorchFix only looks at files that have a `torch` string.
         if m.matches(
             node,
             m.Assign(
@@ -28,7 +26,6 @@ def visit_Assign(self, node):
                         target=m.Attribute(attr=m.Name(value="require_grad"))
                     )
                 ],
-                value=(m.Name("True") | m.Name("False")),
             ),
         ):
             replacement = node.with_deep_changes(