pysyft torrch.jit. script RuntimeError: undefined value _Reduction

Question

I was trying to reproduce Pysyft Asynchronous-federated-learning-on-MNIST from its advanced example. where @torch.jit.script is used before loss function. I am getting this error and have no clue what this is about

RuntimeError: undefined value _Reduction: at /home/ab/.virtualenvs/aic/lib/python3.6/site-packages/syft/generic/frameworks/hook/hook.py:1829:20

reduction = _Reduction.legacy_get_string(size_average, reduce)

It is actually caused by these lines

@torch.jit.script
def loss_fn(pred, target):
    return F.nll_loss(input=pred, target=target)

train_config = sy.TrainConfig(
        model=traced_model,
        loss_fn=loss_fn,
        batch_size=batch_size,
        shuffle=True,
        max_nr_batches=max_nr_batches,
        epochs=1,
        optimizer="SGD",
        optimizer_args={"lr": lr},
    )

Answer 1

Writing answer so it could help others.It turns out that @torch.jit.script needs to be at the top of the file (after import) and I was having it after two function definitions.

Moving it to the top worked