[functorch] fix flake

functorch/op_analysis/gen_data.py

@@ -193,7 +193,10 @@ if True:
     with open('count_ops.txt', 'r') as f:
         opinfo_counts = [i.strip() for i in f.readlines()]
         opinfo_counts = defaultdict(int, {k: v for k, v in zip(opinfo_ops, opinfo_counts)})
-        count_fn = lambda x: opinfo_counts[x['full_name']]
+
+        def count_fn(x):
+            return opinfo_counts[x['full_name']]
+
     with open('run_decompositions.txt', 'r') as f:
         decomposed_ops = [remove_suffix(i.strip(), '.default') for i in f.readlines()]
     gen_data([full_name_check(opinfo_ops), full_name_check(decomposed_ops), count_fn], 'decompositions.txt')

functorch/test/test_ops.py

@@ -13,7 +13,6 @@ from collections import defaultdict
 from contextlib import contextmanager
 from torch.testing._internal.common_device_type import instantiate_device_type_tests
 from torch.testing._internal.common_device_type import ops
-from torch.testing._internal.common_dtype import integral_types
 from torch.testing._internal.common_device_type import \
      toleranceOverride, tol
 from functorch_lagging_op_db import functorch_lagging_op_db