transformer_layer.py: Could not cast value of type NoneType to bool

[WilliamTambellini]

🐛 Bug

Fatal error when trying to torch.jit.save: "Could not cast value of type NoneType to bool"

To Reproduce

modelname='transformer.wmt14.en-fr'
m = torch.hub.load('pytorch/fairseq', modelname, tokenizer='moses', bpe='subword_nmt')
tmodel = m.models[0]
mscript = torch.jit.script(tmodel)
torch.jit.save(mscript, 'tmodel.jit')

Exporting to jit:
Traceback (most recent call last):
  File "testfairseq.py", line 138, in <module>
    mscript = torch.jit.script(tmodel)
  File "/home/wtambellini/.local/lib/python3.8/site-packages/torch/jit/_script.py", line 1265, in script
    return torch.jit._recursive.create_script_module(
  File "/home/wtambellini/.local/lib/python3.8/site-packages/torch/jit/_recursive.py", line 454, in create_script_module
    return create_script_module_impl(nn_module, concrete_type, stubs_fn)
  File "/home/wtambellini/.local/lib/python3.8/site-packages/torch/jit/_recursive.py", line 516, in create_script_module_impl
    script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
  File "/home/wtambellini/.local/lib/python3.8/site-packages/torch/jit/_script.py", line 594, in _construct
    init_fn(script_module)
  File "/home/wtambellini/.local/lib/python3.8/site-packages/torch/jit/_recursive.py", line 494, in init_fn
    scripted = create_script_module_impl(orig_value, sub_concrete_type, stubs_fn)
  File "/home/wtambellini/.local/lib/python3.8/site-packages/torch/jit/_recursive.py", line 516, in create_script_module_impl
    script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
  File "/home/wtambellini/.local/lib/python3.8/site-packages/torch/jit/_script.py", line 594, in _construct
    init_fn(script_module)
  File "/home/wtambellini/.local/lib/python3.8/site-packages/torch/jit/_recursive.py", line 494, in init_fn
    scripted = create_script_module_impl(orig_value, sub_concrete_type, stubs_fn)
  File "/home/wtambellini/.local/lib/python3.8/site-packages/torch/jit/_recursive.py", line 516, in create_script_module_impl
    script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
  File "/home/wtambellini/.local/lib/python3.8/site-packages/torch/jit/_script.py", line 594, in _construct
    init_fn(script_module)
  File "/home/wtambellini/.local/lib/python3.8/site-packages/torch/jit/_recursive.py", line 494, in init_fn
    scripted = create_script_module_impl(orig_value, sub_concrete_type, stubs_fn)
  File "/home/wtambellini/.local/lib/python3.8/site-packages/torch/jit/_recursive.py", line 520, in create_script_module_impl
    create_methods_and_properties_from_stubs(concrete_type, method_stubs, property_stubs)
  File "/home/wtambellini/.local/lib/python3.8/site-packages/torch/jit/_recursive.py", line 371, in create_methods_and_properties_from_stubs
    concrete_type._create_methods_and_properties(property_defs, property_rcbs, method_defs, method_rcbs, method_defaults)
RuntimeError: 
Could not cast value of type NoneType to bool:
  File "/home/wtambellini/repos/fairseq/fairseq/modules/transformer_layer.py", line 502
                incremental_state=incremental_state,
                static_kv=True,
                need_weights=need_attn or (not self.training and self.need_attn),
                                                                                               ~~~~~~~~~~~~~~ <--- HERE

Expected behavior

jit saving to work ?

Environment

• fairseq Version (e.g., 1.0 or main): fairseq version: 1.0.0a0+b5e7b25
• PyTorch Version (e.g., 1.0): torch version: 1.11.0+cu102
• OS (e.g., Linux): Linux
• How you installed fairseq (pip, source): src
• Build command you used (if compiling from source): as explained in the readme
• Python version: 3.8
• CUDA/cuDNN version: 10
• GPU models and configuration: Tesla
• Any other relevant information:

Additional context

JIT scripting cannot handle none type, nothing new. Possible fix:

need_weights=need_attn
if self.need_attn is not None:
   need_weights=need_attn or (not self.training and self.need_attn)

[WilliamTambellini]

@okhonko @stephanpeitz @sarthakgarg what about:

need_weights=need_attn
if self.need_attn is not None and self.training is not None:
   need_weights=need_attn or (not self.training and self.need_attn)

?