Commit fe365d58 authored by Natalia Gimelshein's avatar Natalia Gimelshein
Browse files

don't convert to float bn with affine=False

parent c8bc3e62
......@@ -27,7 +27,7 @@ def BN_convert_float(module):
fn to all modules, parameters, and buffers. Thus we wouldn't
be able to guard the float conversion based on the module type.
'''
if isinstance(module, torch.nn.modules.batchnorm._BatchNorm):
if isinstance(module, torch.nn.modules.batchnorm._BatchNorm) and module.affine is True:
module.float()
for child in module.children():
BN_convert_float(child)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment