[BUG] Cannot save torchscripted model

[dotXem]

Describe the bug
Unable to save a model after successfully torch-scripting it.

To Reproduce

The following snippet:

from effdet import EfficientDet
from effdet.config.model_config import get_efficientdet_config
import torch

architecture = "tf_efficientdet_d1"
config = get_efficientdet_config(architecture)
model = EfficientDet(config, pretrained_backbone=True)
model.toggle_head_bn_level_first()
ts_model = torch.jit.script(model)
torch.jit.save(ts_model, "model.pt")

produces the following output:

(venv) maxime@Morphling: python failing_effdet.py 
Traceback (most recent call last):
  File "failing_effdet.py", line 10, in <module>
    torch.jit.save(ts_model, "model.pt")
  File "/home/maxime/mambaforge/envs/venv/lib/python3.8/site-packages/torch/jit/_serialization.py", line 81, in save
    m.save(f, _extra_files=_extra_files)
  File "/home/maxime/mambaforge/envs/venv/lib/python3.8/site-packages/torch/jit/_script.py", line 714, in save
    return self._c.save(str(f), **kwargs)
RuntimeError: 
Could not export Python function call '_forward'. Remove calls to Python functions before export. Did you forget to add @script or @script_method annotation? If this is a nn.ModuleList, add it to __constants__:
  File "/home/maxime/mambaforge/envs/venv/lib/python3.8/site-packages/effdet/efficientdet.py", line 440
            return self._forward_level_first(x)
        else:
            return self._forward(x)
                   ~~~~~~~~~~~~~ <--- HERE

Expected behavior
Using model.toggle_head_bn_level_first() should have made torchscript happy, because the _forward method, which depends on it, is illegal for torchscript. However, it seems that the function is still checked before saving.

Scripting + saving works fine with torchvision and pytorch lightning models.

Desktop:

• OS: Ubuntu using WSL2 with VSCode
• python==3.8
• effdet==0.3.0
• torch==1.12.0
• CPU only (not attempted on GPU yet)