Support for torch.jit.script

Hey!

Is there any plans to support torch.jit.script for all things in fastai v2 ? I have been trying to export the model generated from https://dev.fast.ai/tutorial.ulmfit witout success.

RuntimeError: Can’t redefine method: forward on class: torch.fastai2.text.models.awdlstm.EmbeddingDropout (addMethod at …/torch/csrc/jit/script/class_type.cpp:73)

Trying also just a few parts, I can see that this fails:

PoolingLinearClassifier(
(layers): Sequential(
(0): LinBnDrop(
(0): BatchNorm1d(1200, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(1): Dropout(p=0.2, inplace=False)
(2): Linear(in_features=1200, out_features=50, bias=False)
(3): ReLU(inplace=True)
)
(1): LinBnDrop(
(0): BatchNorm1d(50, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(1): Dropout(p=0.1, inplace=False)
(2): Linear(in_features=50, out_features=2, bias=False)
)
)
)

With the error:

RuntimeError: Can’t redefine method: forward on class: torch.fastai2.text.models.core.PoolingLinearClassifier (addMethod at …/torch/csrc/jit/script/class_type.cpp:73)

But just trying with the sequential layer, it works.

Will this be supported on the fastai 2 release?

I doubt we’ll see jit support for AWD LSTM for quite a while - I suspect it would need special support from PyTorch to make it work.