torch.nn.utils.fuse_linear_bn_eval

torch.nn.utils.fuse_linear_bn_eval(linear, bn)

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

Parameters

• linear (torch.nn.Linear) – A Linear module.

• bn (torch.nn.modules.batchnorm._BatchNorm) – A BatchNorm module.

Returns

The fused linear module.

Return type

torch.nn.Linear

Note

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