def _test_jit_script(self, tmpdir):
    """Tests that the task can be scripted and saved with JIT then reloaded and used."""
    path = os.path.join(tmpdir, "testing_model.pt")

    model = self.instantiated_task
    model.eval()

    model = torch.jit.script(model)

    torch.jit.save(model, path)
    model = torch.jit.load(path)

    self.check_forward_output(model(self.example_forward_input))
Tests that the task can be scripted and saved with JIT then reloaded and used.
tests/helpers/task_tester.py