inductor: only do the conv+bn folding for the freezing path (#109270)

Pull Request resolved: #109270
Approved by: https://github.com/eellison

torch/_inductor/fx_passes/pre_grad.py

@@ -89,12 +89,11 @@ def fuse_fx(gm: torch.fx.GraphModule, example_inputs):
         gm = permute_matmul_fusion(gm)
 
     # make sure the autograd is disabled.
-    if torch.is_grad_enabled():
+    if torch.is_grad_enabled() or not is_cpu:
         return gm
-    if not is_cpu:
-        return gm
-    gm = remove_identity(gm)
-    gm = fuse_conv_bn(gm)
+    if config.freezing:
+        gm = remove_identity(gm)
+        gm = fuse_conv_bn(gm)
     return gm