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): ...@@ -27,7 +27,7 @@ def BN_convert_float(module):
fn to all modules, parameters, and buffers. Thus we wouldn't fn to all modules, parameters, and buffers. Thus we wouldn't
be able to guard the float conversion based on the module type. 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() module.float()
for child in module.children(): for child in module.children():
BN_convert_float(child) 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