relax assertion on fake shape (#121599)

Summary: Seems like if you use `capture_pre_autograd_graph` fake tensor shapes can be ints instead of symints.

Test Plan: fixes the AssertionError in N5057219

Differential Revision: D54729142

Pull Request resolved: #121599
Approved by: https://github.com/angelayi, https://github.com/BoyuanFeng

Author: avikchaudhuri

torch/fx/experimental/symbolic_shapes.py

@@ -2930,8 +2930,11 @@ def is_dim(src):
             def get_expression(tensor_dim_src):
                 fake = placeholders[source_index[tensor_dim_src.base.name()]]
                 symint = fake.shape[tensor_dim_src.idx]
-                assert isinstance(symint, torch.SymInt)
-                return symint.node.expr
+                if isinstance(symint, torch.SymInt):
+                    return symint.node.expr
+                else:
+                    assert type(symint) is int, f"Expected int, got {type(symint)}"
+                    return symint
 
             for src1, src2 in equalities_inputs.source_pairs:
                 expr1, expr2 = get_expression(src1), get_expression(src2)