torch.nn.utils.fuse_conv_bn_eval — PyTorch 2.7 documentation (original) (raw)

torch.nn.utils.fuse_conv_bn_eval(conv, bn, transpose=False)[source][source]¶

Fuse a convolutional module and a BatchNorm module into a single, new convolutional module.

Parameters

• conv (torch.nn.modules.conv._ConvNd) – A convolutional module.
• bn (torch.nn.modules.batchnorm._BatchNorm) – A BatchNorm module.
• transpose (bool, optional) – If True, transpose the convolutional weight. Defaults to False.

Returns

The fused convolutional module.

Return type

torch.nn.modules.conv._ConvNd

Note

Both conv and bn must be in eval mode, and bn must have its running buffers computed.