Make drop_prob mutable in DropBlock ScriptModules

[mdraw]

This makes it possible to schedule drop_prob during training when used
in JIT mode. The LinearScheduler module still can't be used in JIT
mode yet, but you can change drop_prob from outside of the model, i.e. in
your training loop.

Note that you can still use Python floats in the constructors, but to change the drop_prob buffer values you now need to pass a scalar tensor, e.g.

from dropblock import DropBlock2D
db = DropBlock2D(drop_prob=.5, block_size=3)
m.drop_prob = torch.tensor(.2)

[mdraw]

I've also tried to make LinearScheduler work in JIT (see https://gist.github.com/mdraw/f5d9b24d8c43731756be9245ac4eeffb), but I'm not sure how the error in line 28 can be fixed.

[mdraw]

It looks like buffer mutation is not yet supported at all in ScriptModules, unfortunately: https://github.com/pytorch/pytorch/blob/1409a2afc8d81d911dbcfc2d0b210cf62658237b/test/test_jit.py#L5434.

[miguelvr]

I am questioning the need of having a ScriptModule LinearSchedule, because it is a training tool, not an inference tool.

You don't really want to use the LinearScheduler during inference.