torch.nn.utils.fuse_conv_bn_eval¶

torch.nn.utils.fuse_conv_bn_eval(conv, bn, transpose=False)[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.

Docs