[Dynamo] Fix einops regression (pytorch#154053)

[Dynamo] Fix einops regression (pytorch#153925)

Fixes pytorch#153476

Pull Request resolved: pytorch#153925
Approved by: https://github.com/williamwen42

(cherry picked from commit d44074f)

Co-authored-by: Michael Lazos <[email protected]>

Author: pytorchbot

torch/_dynamo/decorators.py

@@ -685,10 +685,10 @@ def _allow_in_graph_einops():
     if mod is None:
         return
     else:
-        # version > 0.7.0 does allow_in_graph out of tree
+        # version > 0.8.1 does allow_in_graph out of tree
         # for BC we need to keep this in fbcode
         # internal xref https://fb.workplace.com/groups/1026248852325474/permalink/1107135774236781/
-        if Version(mod.__version__) < Version("0.7.0") or is_fbcode():
+        if Version(mod.__version__) <= Version("0.8.1") or is_fbcode():
             import einops
 
             try: