class_docstr: Transformation that fuses `nn.BatchNorm2d` following `nn.Conv2d` into a single `nn.Conv2d`.
  The fusion will be done only if the convolution has the batch normalization as sole following node.
  
  For example, fusion will not be done in the case
  ```
       Conv2d
       /   \
      /     \
  ReLU   BatchNorm2d
  ```
  
  Example:
  ```python
  >>> from transformers.utils.fx import symbolic_trace
  >>> from transformers import AutoModelForImageClassification
  
  >>> from optimum.fx.optimization import FuseBatchNorm2dInConv2d
  
  >>> model = AutoModelForImageClassification.from_pretrained("microsoft/resnet-50")
  >>> model.eval()  # doctest: +IGNORE_RESULT
  
  >>> traced_model = symbolic_trace(
  ...     model,
  ...     input_names=["pixel_values"],
  ...     disable_check=True
  ... )
  
  >>> transformation = FuseBatchNorm2dInConv2d()
  >>> transformed_model = transformation(traced_model)
  ```
class_docstr_tok: ["Transformation","that","fuses","`","nn.BatchNorm2d","`","following","`","nn.Conv2d","`","into","a","single","`","nn.Conv2d","`",".","The","fusion","will","be","done","only","if","the","convolution","has","the","batch","normalization","as","sole","following","node",".","For","example",",","fusion","will","not","be","done","in","the","case","``","`","Conv2d","\/","\\","\/","\\","ReLU","BatchNorm2d","``","`","Example",":","``","`","python",">",">",">","from","transformers.utils.fx","import","symbolic_trace",">",">",">","from","transformers","import","AutoModelForImageClassification",">",">",">","from","optimum.fx.optimization","import","FuseBatchNorm2dInConv2d",">",">",">","model","=","AutoModelForImageClassification.from_pretrained","(","``","microsoft\/resnet-50","''",")",">",">",">","model.eval","(",")","#","doctest",":","+IGNORE_RESULT",">",">",">","traced_model","=","symbolic_trace","(","...","model",",","...","input_names=","[","``","pixel_values","''","]",",","...","disable_check=True","...",")",">",">",">","transformation","=","FuseBatchNorm2dInConv2d","(",")",">",">",">","transformed_model","=","transformation","(","traced_model",")","``","`"]
code: def fuse(self, conv2d: torch.nn.Conv2d, bn2d: torch.nn.BatchNorm2d):
      # handle the case where there is no bias in the conv or the batchnorm has no learnable parameters
      conv_b = (
          conv2d.bias
          if conv2d.bias is not None
          else torch.zeros_like(bn2d.running_mean)
      )
      bn_w = (
          bn2d.weight
          if bn2d.weight is not None
          else torch.ones_like(bn2d.running_mean)
      )
      bn_b = (
          bn2d.bias
          if bn2d.bias is not None
          else torch.ones_like(bn2d.running_mean)
      )
  
      bn_var_rsqrt = torch.rsqrt(bn2d.running_var + bn2d.eps)
  
      conv2d.weight = torch.nn.Parameter(
          conv2d.weight
          * (bn_w * bn_var_rsqrt).reshape(
              [-1] + [1] * (len(conv2d.weight.shape) - 1)
          )
      )
  
      conv2d.bias = torch.nn.Parameter(
          conv_b - bn2d.running_mean * bn_var_rsqrt * bn_w + bn_b
      )
  
      return conv2d
  
imports: import collections
  import copy
  import functools
  import itertools
  import operator
  import warnings
  from abc import ABC, abstractmethod
  from typing import TYPE_CHECKING, List
  import torch
  from transformers.file_utils import add_end_docstrings
  from transformers.utils.fx import _gen_constructor_wrapper
label_desc: Use image node_id 2 for calling the FuseBatchNorm2dInConv2d obj's underlying member method code with example usage: obj.fuse(conv2d, bn2d) and returns: conv2d