Model to TorchScript - Expected a value of type 'float' for argument 'p' but instead found type 'int'

Beginners
1

When i try to convert a model to TorchScript i get - Expected a value of type ‘float’ for argument ‘p’ but instead found type ‘int’

model = AutoModelForQuestionAnswering.from_pretrained("mfeb/albert-xxlarge-v2-squad2")
model.eval().cuda()
scripted_model = torch.jit.script(model)

Stack trace:

Traceback (most recent call last):
  File "/pycharm-2021.2.1/plugins/python/helpers/pydev/pydevd.py", line 1483, in _exec
    pydev_imports.execfile(file, globals, locals)  # execute the script
  File "/pycharm-2021.2.1/plugins/python/helpers/pydev/_pydev_imps/_pydev_execfile.py", line 18, in execfile
    exec(compile(contents+"\n", file, 'exec'), glob, loc)
  File "/Training_util.py", line 118, in <module>
    scripted_model = torch.jit.script(model_en)
  File "/anaconda3/envs/cm_nlp/lib/python3.7/site-packages/torch/jit/_script.py", line 1097, in script
    obj, torch.jit._recursive.infer_methods_to_compile
  File "/anaconda3/envs/cm_nlp/lib/python3.7/site-packages/torch/jit/_recursive.py", line 412, in create_script_module
    return create_script_module_impl(nn_module, concrete_type, stubs_fn)
  File "/anaconda3/envs/cm_nlp/lib/python3.7/site-packages/torch/jit/_recursive.py", line 474, in create_script_module_impl
    script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
  File "/anaconda3/envs/cm_nlp/lib/python3.7/site-packages/torch/jit/_script.py", line 497, in _construct
    init_fn(script_module)
  File "/anaconda3/envs/cm_nlp/lib/python3.7/site-packages/torch/jit/_recursive.py", line 452, in init_fn
    scripted = create_script_module_impl(orig_value, sub_concrete_type, stubs_fn)
  File "/anaconda3/envs/cm_nlp/lib/python3.7/site-packages/torch/jit/_recursive.py", line 474, in create_script_module_impl
    script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
  File "/anaconda3/envs/cm_nlp/lib/python3.7/site-packages/torch/jit/_script.py", line 497, in _construct
    init_fn(script_module)
  File "/anaconda3/envs/cm_nlp/lib/python3.7/site-packages/torch/jit/_recursive.py", line 452, in init_fn
    scripted = create_script_module_impl(orig_value, sub_concrete_type, stubs_fn)
  File "/anaconda3/envs/cm_nlp/lib/python3.7/site-packages/torch/jit/_recursive.py", line 474, in create_script_module_impl
    script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
  File "/anaconda3/envs/cm_nlp/lib/python3.7/site-packages/torch/jit/_script.py", line 497, in _construct
    init_fn(script_module)
  File "/anaconda3/envs/cm_nlp/lib/python3.7/site-packages/torch/jit/_recursive.py", line 452, in init_fn
    scripted = create_script_module_impl(orig_value, sub_concrete_type, stubs_fn)
  File "/anaconda3/envs/cm_nlp/lib/python3.7/site-packages/torch/jit/_recursive.py", line 478, in create_script_module_impl
    create_methods_and_properties_from_stubs(concrete_type, method_stubs, property_stubs)
  File "/anaconda3/envs/cm_nlp/lib/python3.7/site-packages/torch/jit/_recursive.py", line 355, in create_methods_and_properties_from_stubs
    concrete_type._create_methods_and_properties(property_defs, property_rcbs, method_defs, method_rcbs, method_defaults)
RuntimeError: 
dropout(Tensor input, float p=0.5, bool training=True, bool inplace=False) -> (Tensor):
Expected a value of type 'float' for argument 'p' but instead found type 'int'.
:
  File "/anaconda3/envs/cm_nlp/lib/python3.7/site-packages/torch/nn/modules/dropout.py", line 58
    def forward(self, input: Tensor) -> Tensor:
        return F.dropout(input, self.p, self.training, self.inplace)
               ~~~~~~~~~ <--- HERE

Can anyone guide me where is my mistake?